From 74338e94f10ec1f2cc9f3f1b5060612082f45c3d Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Wed, 22 Jul 2026 21:15:03 +0800 Subject: [PATCH] [Kernel] Phase 4 batch-3: migrate tangled JIT subsystems + new groups into kernels.ops (RFC #29630) (#32045) Co-authored-by: Claude Opus 4.8 (1M context) --- .../jit_kernel/benchmark/kv_canary/utils.py | 2 +- python/sglang/jit_kernel/dsa/__init__.py | 30 - python/sglang/jit_kernel/dsv4/__init__.py | 63 - .../jit_kernel/kv_canary/plan/__init__.py | 1 - .../sglang/jit_kernel/minimax_m3/__init__.py | 25 - .../jit_kernel/tests/deepseek_v4/common.py | 5 +- .../tests/kv_canary/_canary_helpers.py | 8 +- .../jit_kernel/tests/kv_canary/_constants.py | 2 +- .../tests/kv_canary/_differential.py | 26 +- .../jit_kernel/tests/kv_canary/_fixtures.py | 6 +- .../tests/kv_canary/_hand_oracle.py | 2 +- .../jit_kernel/tests/kv_canary/_invariants.py | 6 +- .../jit_kernel/tests/test_minimax_m3_mxfp8.py | 2 +- .../tests/test_minimax_m3_rmsnorm.py | 2 +- .../jit_kernel/tests/test_moe_topk_sigmoid.py | 2 +- .../jit_kernel/trtllm_lora_temp/__init__.py | 19 - .../data/csrc/fused_activation_quant.cuh | 246 - .../data/csrc/fused_permute_quant.cuh | 307 -- .../csrc/trtllm_fused_moe_kernel_launcher.cu | 4380 ---------------- .../data/csrc/trtllm_fused_moe_runner.cu | 1116 ----- .../flashinfer/trtllm/fused_moe/DevKernel.h | 472 -- .../flashinfer/trtllm/fused_moe/runner.h | 586 --- python/sglang/kernels/ops/__init__.py | 3 + .../ops/attention}/add_constant.py | 0 .../ops/attention}/clamp_position.py | 0 .../ops/attention}/concat_mla.py | 0 .../cutedsl_fp8_paged_mqa_logits.py | 0 .../ops/attention}/cutedsl_gdn.py | 0 .../ops/attention}/cutedsl_kda.py | 0 .../kernels/ops/attention/dsa/__init__.py | 32 + .../dsa/cutedsl_paged_mqa_logits.py | 2 +- .../ops/attention}/dsa/paged_mqa_logits.py | 2 +- .../kernels/ops/attention/dsv4/__init__.py | 65 + .../ops/attention}/dsv4/attn.py | 0 .../ops/attention}/dsv4/c128_cleanup.py | 0 .../ops/attention}/dsv4/compress.py | 0 .../ops/attention/dsv4/compress_c128_hip.py | 4 +- .../ops/attention}/dsv4/compress_old.py | 0 .../ops/attention}/dsv4/elementwise.py | 0 .../ops/attention}/dsv4/fp8_wo_a.py | 0 .../attention/dsv4/fused_compress_triton.py | 2 +- .../ops/attention}/dsv4/gemm.py | 0 .../ops/attention}/dsv4/moe.py | 2 +- .../ops/attention}/dsv4/online_c128_mtp.py | 2 +- .../ops/attention}/dsv4/topk.py | 0 .../ops/attention}/dsv4/utils.py | 0 .../ops/attention}/fixup_zero_kv.py | 0 .../ops/attention}/flash_attention.py | 0 .../ops/attention}/flash_attention_v3.py | 0 .../ops/attention}/flash_attention_v4.py | 2 +- .../ops/attention}/flash_attn/__init__.py | 0 .../ops/attention}/flash_attn/cute/.flake8 | 0 .../ops/attention}/flash_attn/cute/AUTHORS | 0 .../ops/attention}/flash_attn/cute/LICENSE | 0 .../attention}/flash_attn/cute/MANIFEST.in | 0 .../ops/attention}/flash_attn/cute/README.md | 0 .../attention}/flash_attn/cute/__init__.py | 0 .../flash_attn/cute/ampere_helpers.py | 0 .../ops/attention}/flash_attn/cute/barrier.py | 0 .../attention}/flash_attn/cute/bench_utils.py | 0 .../attention}/flash_attn/cute/benchmark.py | 0 .../cute/benchmark_flash_attention_fp8.py | 4 +- .../flash_attn/cute/blackwell_helpers.py | 2 +- .../attention}/flash_attn/cute/block_info.py | 5 +- .../flash_attn/cute/block_sparse_utils.py | 10 +- .../flash_attn/cute/block_sparsity.py | 2 +- .../attention}/flash_attn/cute/cache_utils.py | 2 +- .../flash_attn/cute/compute_block_sparsity.py | 14 +- .../attention}/flash_attn/cute/copy_utils.py | 0 .../flash_attn/cute/cu_blocks_kernels.py | 0 .../flash_attn/cute/cute_dsl_ptxas.py | 0 .../flash_attn/cute/cute_dsl_utils.py | 0 .../attention}/flash_attn/cute/fa_logging.py | 0 .../attention}/flash_attn/cute/fast_math.py | 0 .../attention}/flash_attn/cute/flash_fwd.py | 33 +- .../flash_attn/cute/flash_fwd_combine.py | 8 +- .../flash_attn/cute/flash_fwd_mla_sm100.py | 36 +- .../flash_attn/cute/flash_fwd_sm100.py | 47 +- .../flash_attn/cute/flash_fwd_sm120.py | 4 +- .../flash_attn/cute/flash_fwd_sm90.py | 39 +- .../attention}/flash_attn/cute/interface.py | 42 +- .../ops/attention}/flash_attn/cute/mask.py | 8 +- .../flash_attn/cute/mma_sm100_desc.py | 0 .../flash_attn/cute/named_barrier.py | 0 .../attention}/flash_attn/cute/pack_gqa.py | 2 +- .../attention}/flash_attn/cute/paged_kv.py | 2 +- .../attention}/flash_attn/cute/pipeline.py | 0 .../attention}/flash_attn/cute/pyproject.toml | 0 .../attention}/flash_attn/cute/seqlen_info.py | 0 .../flash_attn/cute/shearing_bias.py | 16 +- .../cute/sm100_hd256_2cta_fmha_forward.py | 24 +- .../flash_attn/cute/sm90_config_search.py | 0 .../ops/attention}/flash_attn/cute/softmax.py | 6 +- .../ops/attention}/flash_attn/cute/testing.py | 0 .../flash_attn/cute/tile_scheduler.py | 4 +- .../flash_attn/cute/topk_gather_kv.py | 4 +- .../ops/attention}/flash_attn/cute/utils.py | 0 .../ops/attention}/fused_metadata_copy.py | 0 .../ops/attention}/fused_qknorm_rope.py | 0 .../ops/attention}/fused_store_index_cache.py | 0 .../ops/attention}/hadamard.py | 0 .../kernels/ops/attention/log_scaling_tau.py | 2 +- .../ops/attention}/minimax_decode_topk.py | 0 .../ops/attention/minimax_m3_qk_norm_rope.py} | 0 .../ops/attention}/minimax_qknorm_rope.py | 0 .../decode/flash_with_topk_idx.py | 4 +- .../attention}/mla_kv_pack_quantize_fp8.py | 0 .../ops/attention}/rope.py | 0 .../sparse_mla_q8kv8_prefill_sm90.py | 0 .../ops/attention/triton_gdn_fused_proj.py} | 0 python/sglang/kernels/ops/attention/utils.py | 2 +- .../sglang/kernels/ops/diffusion/__init__.py | 8 +- .../ops}/diffusion/causal_conv3d_cat_pad.py | 0 .../diffusion/cutedsl/common/norm_fusion.py | 2 +- .../ops}/diffusion/cutedsl/common/reduce.py | 0 .../cutedsl/norm_tanh_mul_add_norm_scale.py | 4 +- .../scale_residual_norm_scale_shift.py | 8 +- .../ops}/diffusion/cutedsl/utils.py | 0 .../diffusion/flydsl/fused_residual_norm.py | 0 .../ops}/diffusion/group_norm_silu.py | 2 +- .../ops}/diffusion/ltx2_qknorm_split_rope.py | 0 .../ops}/diffusion/norm_scale_shift_native.py | 0 .../ops}/diffusion/qknorm_rope.py | 0 .../ops}/diffusion/residual_gate_add.py | 0 .../ops/diffusion}/timestep_embedding.py | 0 .../diffusion/triton/causal_conv3d_pad.py | 0 .../ops}/diffusion/triton/group_norm_silu.py | 0 .../ops}/diffusion/triton/ltx2_ada_values.py | 0 .../ops}/diffusion/triton/ltx2_rotary.py | 0 .../ops}/diffusion/triton/mps_fallback.py | 0 .../ops}/diffusion/triton/norm.py | 0 .../ops}/diffusion/triton/npu_fallback.py | 0 .../ops}/diffusion/triton/rmsnorm_onepass.py | 0 .../ops}/diffusion/triton/rotary.py | 0 .../ops}/diffusion/triton/sana_wm_gdn.py | 2 +- .../diffusion/triton/sana_wm_gdn_chunkwise.py | 2 +- .../ops}/diffusion/triton/scale_shift.py | 0 .../ops}/diffusion/triton/torch_fallback.py | 0 .../ops}/diffusion/triton/varlen_pack_pad.py | 0 .../diffusion/triton/zimage_native_norm.py | 0 .../ops}/kv_canary/__init__.py | 0 .../ops}/kv_canary/consts.py | 0 .../kernels/ops/kv_canary/plan/__init__.py | 1 + .../ops}/kv_canary/plan/api.py | 10 +- .../ops}/kv_canary/plan/entries_kernel.py | 0 .../ops}/kv_canary/plan/offsets_kernel.py | 6 +- .../ops}/kv_canary/plan/utils.py | 0 .../ops}/kv_canary/plan_ref.py | 6 +- .../ops}/kv_canary/scatter_req_token_ids.py | 0 .../ops}/kv_canary/verify.py | 4 +- .../ops}/kv_canary/verify_ref.py | 6 +- .../ops}/kv_canary/write.py | 8 +- .../ops}/kv_canary/write_ref.py | 8 +- .../sglang/kernels/ops/layernorm/__init__.py | 6 +- .../ops/layernorm/minimax_m3_rmsnorm.py} | 0 .../ops}/lplb/__init__.py | 0 .../ops}/lplb/cublasdx_solver.py | 2 +- .../ops}/lplb/cuda_solver.py | 0 .../ops}/lplb/shmem_budget.py | 0 .../ops}/lplb/torch_solver.py | 6 +- python/sglang/kernels/ops/model/__init__.py | 1 + .../kernels/ops/model/inkling/__init__.py | 1 + .../ops/model/inkling}/inkling_all_reduce.py | 0 .../ops/model/inkling}/inkling_ar_fused.py | 0 .../inkling}/inkling_ar_scattered_sconv.py | 0 .../model/inkling}/inkling_attn_prologue.py | 0 .../inkling}/inkling_gate_topk_renorm.py | 0 .../ops/model/inkling}/inkling_rel_proj.py | 0 .../ops/model/inkling}/inkling_row_scale.py | 0 python/sglang/kernels/ops/moe/__init__.py | 2 +- .../sglang/kernels/ops/moe/ep_moe_kernels.py | 2 +- .../ops/moe}/kpool_topk_transform.py | 0 .../ops/moe/minimax_m3_swiglu.py} | 0 .../ops/moe}/moe_align.py | 0 .../ops/moe}/moe_finalize_fuse_shared.py | 0 .../ops/moe}/moe_fused_gate.py | 0 .../ops/moe}/moe_lora_align.py | 0 .../ops/moe}/moe_permute_prepare.py | 0 .../ops/moe}/moe_topk_sigmoid.py | 0 .../ops/moe}/moe_wna16_marlin.py | 0 .../kernels/ops/moe/mxfp8_moe_amd_gfx95.py | 5 +- .../ops/moe/triton_hash_topk.py} | 0 .../ops/moe/triton_sigmoid_gate_mul.py} | 0 .../ops/moe}/trtllm_lora_temp/SOURCE.md | 0 .../ops/moe/trtllm_lora_temp/__init__.py | 21 + .../ops/moe}/trtllm_lora_temp/core.py | 4 +- .../data/csrc/fused_activation_quant.cuh | 259 + .../trtllm_fused_moe_dev_kernel.cu | 639 ++- .../data/csrc/fused_permute_quant.cuh | 322 ++ .../csrc/trtllm_fused_moe_kernel_launcher.cu | 4383 +++++++++++++++++ .../data/csrc/trtllm_fused_moe_runner.cu | 927 ++++ .../flashinfer/trtllm/fused_moe/DevKernel.h | 481 ++ .../flashinfer/trtllm/fused_moe/runner.h | 540 ++ .../ops/moe}/trtllm_lora_temp/jit.py | 0 .../kimi_k2_moe_fused_gate.py | 0 .../trtllm_lora_temp/moe_lora_merged_align.py | 0 .../ops/moe}/trtllm_lora_temp/topk_pack.py | 0 .../trtllm_lora_temp/topk_softmax_pack.py | 0 .../moe/trtllm_lora_temp/virtual_experts.py | 6 +- .../sglang/kernels/ops/moe/virtual_experts.py | 4 +- .../ops/quantization}/dsv32/__init__.py | 0 .../ops/quantization}/dsv32/elementwise.py | 0 .../layers/attention/backends/flash_attn.py | 2 +- .../runtime/layers/attention/layer.py | 4 +- .../runtime/layers/elementwise.py | 2 +- .../runtime/layers/fused_scale_shift_gate.py | 2 +- .../runtime/layers/layernorm.py | 22 +- .../runtime/layers/parallel_conv.py | 4 +- .../runtime/layers/rotary_embedding/utils.py | 2 +- .../runtime/layers/visual_embedding.py | 2 +- .../runtime/models/dits/krea2.py | 6 +- .../runtime/models/dits/ltx_2.py | 8 +- .../runtime/models/dits/sana_wm_components.py | 4 +- .../runtime/models/dits/zimage.py | 4 +- .../models/upsampler/latent_upsampler.py | 2 +- .../runtime/models/vaes/hunyuanvae.py | 2 +- .../multimodal_gen/runtime/platforms/rocm.py | 4 +- .../test_latent_upsampler_group_norm_silu.py | 4 +- .../srt/eplb/expert_location_dispatch.py | 2 +- python/sglang/srt/eplb/lplb_solver.py | 4 +- python/sglang/srt/kv_canary/buffer_group.py | 2 +- python/sglang/srt/kv_canary/config.py | 2 +- python/sglang/srt/kv_canary/endpoint.py | 6 +- python/sglang/srt/kv_canary/perturb/utils.py | 2 +- .../kv_canary/pool_patcher/buffer_alloc.py | 4 +- .../req_to_expected_token_ids_manager.py | 2 +- .../srt/kv_canary/runner/canary_manager.py | 2 +- .../srt/kv_canary/runner/health_checker.py | 2 +- .../srt/kv_canary/runner/kernel_launcher.py | 8 +- .../srt/kv_canary/runner/stats_logger.py | 2 +- .../srt/kv_canary/runner/swa_divergence.py | 2 +- .../kv_canary/runner/violation_reporter.py | 4 +- .../single_forward_manager/manager.py | 4 +- python/sglang/srt/kv_canary/state.py | 4 +- .../srt/kv_canary/sweep_plan_builder.py | 2 +- .../layers/attention/deepseek_v4_backend.py | 2 +- .../srt/layers/attention/dsa/dsa_indexer.py | 12 +- .../layers/attention/dsa/dsa_topk_backend.py | 2 +- .../srt/layers/attention/dsa_backend.py | 12 +- .../srt/layers/attention/dsv4/compressor.py | 9 +- .../layers/attention/dsv4/compressor_v2.py | 4 +- .../srt/layers/attention/dsv4/indexer.py | 2 +- .../srt/layers/attention/dsv4/metadata.py | 6 +- .../dual_chunk_flashattention_backend.py | 2 +- .../attention/flashattention_backend.py | 4 +- .../layers/attention/linear/gdn_backend.py | 4 +- .../attention/linear/kernels/gdn_cutedsl.py | 4 +- .../attention/linear/kernels/kda_cutedsl.py | 4 +- .../attention/tokenspeed_mla_backend.py | 4 +- .../layers/attention/trtllm_mla_backend.py | 2 +- python/sglang/srt/layers/attention/vision.py | 4 +- python/sglang/srt/layers/layernorm.py | 4 +- .../moe/fused_moe_triton/fused_marlin_moe.py | 2 +- python/sglang/srt/layers/moe/hash_topk.py | 2 +- python/sglang/srt/layers/moe/mega_moe.py | 2 +- .../srt/layers/moe/moe_runner/deep_gemm.py | 4 +- .../moe/moe_runner/flashinfer_trtllm.py | 2 +- .../moe/moe_runner/triton_utils/fused_moe.py | 2 +- .../triton_utils/moe_align_block_size.py | 2 +- .../triton_utils/sigmoid_gate_topk_renorm.py | 4 +- python/sglang/srt/layers/moe/topk.py | 26 +- .../srt/layers/rotary_embedding/base.py | 10 +- .../srt/layers/rotary_embedding/mrope.py | 2 +- .../sglang/srt/lora/lora_moe_runner_marlin.py | 2 +- python/sglang/srt/lora/lora_moe_runners.py | 2 +- .../srt/lora/marlin_lora_temp/moe_runner.py | 2 +- .../experimental_sgl_trtllm_moe.py | 4 +- .../lora/trtllm_lora_temp/lora_dispatch.py | 12 +- .../srt/lora/trtllm_lora_temp/lora_layer.py | 2 +- .../srt/lora/trtllm_lora_temp/moe_overlap.py | 12 +- .../srt/lora/trtllm_lora_temp/sgl_fp8_moe.py | 2 +- .../srt/mem_cache/deepseek_v4_memory_pool.py | 4 +- .../srt/model_executor/forward_batch_info.py | 2 +- .../attention_forward_methods/forward_mha.py | 2 +- python/sglang/srt/models/deepseek_v2.py | 9 +- python/sglang/srt/models/deepseek_v4.py | 4 +- .../sglang/srt/models/deepseek_v4_dspark.py | 2 +- python/sglang/srt/models/inkling.py | 6 +- .../sglang/srt/models/inkling_common/attn.py | 18 +- .../sglang/srt/models/inkling_common/moe.py | 4 +- python/sglang/srt/models/longcat_flash.py | 2 +- python/sglang/srt/models/mellum.py | 2 +- python/sglang/srt/models/minimax_m3.py | 10 +- python/sglang/srt/models/qwen2_moe.py | 2 +- python/sglang/srt/models/qwen3_5.py | 7 +- python/sglang/srt/models/qwen3_moe.py | 2 +- python/sglang/srt/models/qwen3_next.py | 8 +- python/sglang/srt/models/sarvam_moe.py | 2 +- python/sglang/srt/models/utils.py | 2 +- python/sglang/srt/server_args.py | 2 +- .../attention_methods/dsa_attention.py | 7 +- python/sglang/test/kv_canary/fixtures.py | 4 +- .../sglang/test/kv_canary/runner_test_base.py | 4 +- .../dual_chunk/test_dual_chunk_flash_attn.py | 4 +- test/registered/eplb/test_lplb_distributed.py | 6 +- .../gemm/test_linear_bf16_fp32_hpc.py | 4 +- .../jit/benchmark/bench_add_constant.py | 5 +- .../jit/benchmark/bench_clamp_position.py | 2 +- .../jit/benchmark/bench_concat_mla.py | 4 +- .../jit/benchmark/bench_fused_qknorm_rope.py | 2 +- .../jit/benchmark/bench_hadamard.py | 2 +- .../bench_mla_kv_pack_quantize_fp8.py | 4 +- .../jit/benchmark/bench_moe_fused_gate.py | 2 +- .../jit/benchmark/bench_online_c128_mtp.py | 4 +- test/registered/jit/benchmark/bench_rope.py | 8 +- .../bench_sparse_mla_q8kv8_prefill_sm90.py | 4 +- test/registered/jit/benchmark/bench_topk.py | 2 +- .../diffusion/bench_causal_conv3d_cat_pad.py | 4 +- .../diffusion/bench_group_norm_silu.py | 2 +- .../diffusion/bench_ltx2_qknorm_split_rope.py | 2 +- .../benchmark/diffusion/bench_norm_impls.py | 4 +- .../benchmark/diffusion/bench_qknorm_rope.py | 2 +- .../diffusion/bench_qwen_image_modulation.py | 4 +- .../diffusion/bench_residual_gate_add.py | 4 +- .../jit/benchmark/kv_canary/bench_plan.py | 6 +- .../kv_canary/bench_scatter_req_token_ids.py | 2 +- .../jit/benchmark/kv_canary/bench_verify.py | 4 +- .../jit/benchmark/kv_canary/bench_write.py | 6 +- .../minimax/bench_minimax_decode_topk.py | 2 +- .../minimax/bench_minimax_qknorm_rope.py | 2 +- .../jit/deepseek_v4/test_c128_v2.py | 2 +- test/registered/jit/deepseek_v4/test_c4_v2.py | 2 +- .../jit/deepseek_v4/test_fp4_indexer.py | 12 +- .../jit/deepseek_v4/test_fp8_wo_a.py | 4 +- .../jit/deepseek_v4/test_topk_v2.py | 2 +- .../diffusion/test_causal_conv3d_cat_pad.py | 6 +- .../jit/diffusion/test_flydsl_fused_norm.py | 4 +- .../diffusion/test_fused_norm_scale_shift.py | 2 +- .../jit/diffusion/test_group_norm_silu.py | 4 +- .../jit/diffusion/test_ltx2_ada_values.py | 2 +- .../diffusion/test_ltx2_qknorm_split_rope.py | 2 +- .../jit/diffusion/test_qknorm_rope.py | 2 +- .../diffusion/test_qwen_image_modulation.py | 6 +- .../jit/diffusion/test_residual_gate_add.py | 2 +- .../jit/diffusion/test_varlen_pack_pad.py | 4 +- .../test_varlen_uspattn_equivalence.py | 6 +- .../jit/kv_canary/test_const_sync.py | 2 +- .../jit/kv_canary/test_kernel_config.py | 16 +- .../jit/kv_canary/test_pipeline_e2e.py | 38 +- .../jit/kv_canary/test_plan_hand.py | 12 +- .../kv_canary/test_scatter_req_token_ids.py | 2 +- .../jit/kv_canary/test_verify_fuzz.py | 12 +- .../jit/kv_canary/test_verify_hand.py | 32 +- .../jit/kv_canary/test_write_fuzz.py | 12 +- .../jit/kv_canary/test_write_hand.py | 26 +- .../jit/minimax/test_minimax_decode_topk.py | 2 +- .../test_minimax_decode_topk_page_table.py | 2 +- .../jit/minimax/test_minimax_qknorm_rope.py | 2 +- test/registered/jit/test_add_constant.py | 2 +- test/registered/jit/test_clamp_position.py | 2 +- test/registered/jit/test_concat_mla.py | 4 +- ...epseek_v4_compress_state_runtime_shapes.py | 10 +- .../jit/test_dsv32_indexer_fusion.py | 12 +- .../jit/test_flash_attention_3_only_qv.py | 2 +- test/registered/jit/test_flash_attention_4.py | 2 +- .../jit/test_fused_metadata_copy.py | 24 +- .../jit/test_fused_store_index_cache.py | 2 +- test/registered/jit/test_hadamard_jit.py | 2 +- .../jit/test_inkling_attn_prologue_tau.py | 4 +- test/registered/jit/test_inkling_rel_proj.py | 2 +- test/registered/jit/test_inkling_row_scale.py | 4 +- .../jit/test_kpool_topk_transform.py | 2 +- .../jit/test_minimax_m3_qk_norm_rope.py | 2 +- .../jit/test_mla_kv_pack_quantize_fp8.py | 4 +- .../jit/test_moe_align_block_size.py | 2 +- test/registered/jit/test_moe_fused_gate.py | 2 +- .../jit/test_moe_lora_align_block_size.py | 2 +- test/registered/jit/test_moe_wna16_marlin.py | 2 +- test/registered/jit/test_pos_enc.py | 2 +- test/registered/jit/test_rope.py | 4 +- test/registered/jit/test_sigmoid_gate_mul.py | 8 +- .../jit/test_sparse_mla_q8kv8_prefill_sm90.py | 16 +- .../registered/jit/test_timestep_embedding.py | 4 +- .../test_cute_dsl_fp8_paged_mqa_logits.py | 2 +- .../kernels/test_deepgemm_paged_mqa_logits.py | 2 +- test/registered/kernels/test_dsa_indexer.py | 2 +- .../kernels/test_sigmoid_gate_mul.py | 10 +- .../kv_canary/test_self_unit_buffer_alloc.py | 2 +- .../kv_canary/test_self_unit_endpoint.py | 6 +- .../kv_canary/test_self_unit_perturb.py | 2 +- .../kv_canary/test_self_unit_pool_patcher.py | 4 +- .../kv_canary/test_self_unit_runner_health.py | 2 +- .../test_self_unit_runner_per_forward.py | 6 +- .../test_self_unit_runner_swa_divergence.py | 2 +- .../kv_canary/test_self_unit_violation.py | 6 +- .../lora/test_fused_moe_lora_kernel.py | 5 +- .../mock_model/test_self_unit_oracle.py | 2 +- .../test_self_unit_oracle_torch_vs_ref.py | 2 +- .../test_experimental_sgl_marlin_alignment.py | 6 +- 389 files changed, 8186 insertions(+), 8158 deletions(-) delete mode 100644 python/sglang/jit_kernel/dsa/__init__.py delete mode 100644 python/sglang/jit_kernel/dsv4/__init__.py delete mode 100644 python/sglang/jit_kernel/kv_canary/plan/__init__.py delete mode 100644 python/sglang/jit_kernel/minimax_m3/__init__.py delete mode 100644 python/sglang/jit_kernel/trtllm_lora_temp/__init__.py delete mode 100644 python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_activation_quant.cuh delete mode 100644 python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_permute_quant.cuh delete mode 100644 python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/trtllm_fused_moe_kernel_launcher.cu delete mode 100644 python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/trtllm_fused_moe_runner.cu delete mode 100644 python/sglang/jit_kernel/trtllm_lora_temp/data/include/flashinfer/trtllm/fused_moe/DevKernel.h delete mode 100644 python/sglang/jit_kernel/trtllm_lora_temp/data/include/flashinfer/trtllm/fused_moe/runner.h rename python/sglang/{jit_kernel => kernels/ops/attention}/add_constant.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/clamp_position.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/concat_mla.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/cutedsl_fp8_paged_mqa_logits.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/cutedsl_gdn.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/cutedsl_kda.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/dsa/cutedsl_paged_mqa_logits.py (99%) rename python/sglang/{jit_kernel => kernels/ops/attention}/dsa/paged_mqa_logits.py (98%) rename python/sglang/{jit_kernel => kernels/ops/attention}/dsv4/attn.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/dsv4/c128_cleanup.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/dsv4/compress.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/dsv4/compress_old.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/dsv4/elementwise.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/dsv4/fp8_wo_a.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/dsv4/gemm.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/dsv4/moe.py (98%) rename python/sglang/{jit_kernel => kernels/ops/attention}/dsv4/online_c128_mtp.py (99%) rename python/sglang/{jit_kernel => kernels/ops/attention}/dsv4/topk.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/dsv4/utils.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/fixup_zero_kv.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attention.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attention_v3.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attention_v4.py (99%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/__init__.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/.flake8 (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/AUTHORS (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/LICENSE (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/MANIFEST.in (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/README.md (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/__init__.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/ampere_helpers.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/barrier.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/bench_utils.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/benchmark.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/benchmark_flash_attention_fp8.py (99%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/blackwell_helpers.py (99%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/block_info.py (98%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/block_sparse_utils.py (99%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/block_sparsity.py (99%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/cache_utils.py (99%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/compute_block_sparsity.py (97%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/copy_utils.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/cu_blocks_kernels.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/cute_dsl_ptxas.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/cute_dsl_utils.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/fa_logging.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/fast_math.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/flash_fwd.py (98%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/flash_fwd_combine.py (99%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/flash_fwd_mla_sm100.py (99%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/flash_fwd_sm100.py (99%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/flash_fwd_sm120.py (95%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/flash_fwd_sm90.py (98%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/interface.py (98%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/mask.py (99%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/mma_sm100_desc.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/named_barrier.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/pack_gqa.py (99%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/paged_kv.py (99%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/pipeline.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/pyproject.toml (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/seqlen_info.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/shearing_bias.py (97%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/sm100_hd256_2cta_fmha_forward.py (99%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/sm90_config_search.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/softmax.py (99%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/testing.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/tile_scheduler.py (99%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/topk_gather_kv.py (98%) rename python/sglang/{jit_kernel => kernels/ops/attention}/flash_attn/cute/utils.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/fused_metadata_copy.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/fused_qknorm_rope.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/fused_store_index_cache.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/hadamard.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/minimax_decode_topk.py (100%) rename python/sglang/{jit_kernel/minimax_m3/qk_norm_rope.py => kernels/ops/attention/minimax_m3_qk_norm_rope.py} (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/minimax_qknorm_rope.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/mla_kv_pack_quantize_fp8.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/rope.py (100%) rename python/sglang/{jit_kernel => kernels/ops/attention}/sparse_mla_q8kv8_prefill_sm90.py (100%) rename python/sglang/{jit_kernel/triton/gdn_fused_proj.py => kernels/ops/attention/triton_gdn_fused_proj.py} (100%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/causal_conv3d_cat_pad.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/cutedsl/common/norm_fusion.py (99%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/cutedsl/common/reduce.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py (98%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/cutedsl/scale_residual_norm_scale_shift.py (98%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/cutedsl/utils.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/flydsl/fused_residual_norm.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/group_norm_silu.py (91%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/ltx2_qknorm_split_rope.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/norm_scale_shift_native.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/qknorm_rope.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/residual_gate_add.py (100%) rename python/sglang/{jit_kernel => kernels/ops/diffusion}/timestep_embedding.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/triton/causal_conv3d_pad.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/triton/group_norm_silu.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/triton/ltx2_ada_values.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/triton/ltx2_rotary.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/triton/mps_fallback.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/triton/norm.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/triton/npu_fallback.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/triton/rmsnorm_onepass.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/triton/rotary.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/triton/sana_wm_gdn.py (99%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/triton/sana_wm_gdn_chunkwise.py (99%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/triton/scale_shift.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/triton/torch_fallback.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/triton/varlen_pack_pad.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/diffusion/triton/zimage_native_norm.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/kv_canary/__init__.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/kv_canary/consts.py (100%) create mode 100644 python/sglang/kernels/ops/kv_canary/plan/__init__.py rename python/sglang/{jit_kernel => kernels/ops}/kv_canary/plan/api.py (96%) rename python/sglang/{jit_kernel => kernels/ops}/kv_canary/plan/entries_kernel.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/kv_canary/plan/offsets_kernel.py (98%) rename python/sglang/{jit_kernel => kernels/ops}/kv_canary/plan/utils.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/kv_canary/plan_ref.py (98%) rename python/sglang/{jit_kernel => kernels/ops}/kv_canary/scatter_req_token_ids.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/kv_canary/verify.py (99%) rename python/sglang/{jit_kernel => kernels/ops}/kv_canary/verify_ref.py (98%) rename python/sglang/{jit_kernel => kernels/ops}/kv_canary/write.py (98%) rename python/sglang/{jit_kernel => kernels/ops}/kv_canary/write_ref.py (97%) rename python/sglang/{jit_kernel/minimax_m3/rmsnorm.py => kernels/ops/layernorm/minimax_m3_rmsnorm.py} (100%) rename python/sglang/{jit_kernel => kernels/ops}/lplb/__init__.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/lplb/cublasdx_solver.py (85%) rename python/sglang/{jit_kernel => kernels/ops}/lplb/cuda_solver.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/lplb/shmem_budget.py (100%) rename python/sglang/{jit_kernel => kernels/ops}/lplb/torch_solver.py (96%) create mode 100644 python/sglang/kernels/ops/model/__init__.py create mode 100644 python/sglang/kernels/ops/model/inkling/__init__.py rename python/sglang/{jit_kernel => kernels/ops/model/inkling}/inkling_all_reduce.py (100%) rename python/sglang/{jit_kernel => kernels/ops/model/inkling}/inkling_ar_fused.py (100%) rename python/sglang/{jit_kernel => kernels/ops/model/inkling}/inkling_ar_scattered_sconv.py (100%) rename python/sglang/{jit_kernel => kernels/ops/model/inkling}/inkling_attn_prologue.py (100%) rename python/sglang/{jit_kernel => kernels/ops/model/inkling}/inkling_gate_topk_renorm.py (100%) rename python/sglang/{jit_kernel => kernels/ops/model/inkling}/inkling_rel_proj.py (100%) rename python/sglang/{jit_kernel => kernels/ops/model/inkling}/inkling_row_scale.py (100%) rename python/sglang/{jit_kernel => kernels/ops/moe}/kpool_topk_transform.py (100%) rename python/sglang/{jit_kernel/minimax_m3/swiglu.py => kernels/ops/moe/minimax_m3_swiglu.py} (100%) rename python/sglang/{jit_kernel => kernels/ops/moe}/moe_align.py (100%) rename python/sglang/{jit_kernel => kernels/ops/moe}/moe_finalize_fuse_shared.py (100%) rename python/sglang/{jit_kernel => kernels/ops/moe}/moe_fused_gate.py (100%) rename python/sglang/{jit_kernel => kernels/ops/moe}/moe_lora_align.py (100%) rename python/sglang/{jit_kernel => kernels/ops/moe}/moe_permute_prepare.py (100%) rename python/sglang/{jit_kernel => kernels/ops/moe}/moe_topk_sigmoid.py (100%) rename python/sglang/{jit_kernel => kernels/ops/moe}/moe_wna16_marlin.py (100%) rename python/sglang/{jit_kernel/triton/hash_topk.py => kernels/ops/moe/triton_hash_topk.py} (100%) rename python/sglang/{jit_kernel/triton/sigmoid_gate_mul.py => kernels/ops/moe/triton_sigmoid_gate_mul.py} (100%) rename python/sglang/{jit_kernel => kernels/ops/moe}/trtllm_lora_temp/SOURCE.md (100%) rename python/sglang/{jit_kernel => kernels/ops/moe}/trtllm_lora_temp/core.py (99%) create mode 100644 python/sglang/kernels/ops/moe/trtllm_lora_temp/data/csrc/fused_activation_quant.cuh rename python/sglang/{jit_kernel => kernels/ops/moe}/trtllm_lora_temp/data/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_dev_kernel.cu (63%) create mode 100644 python/sglang/kernels/ops/moe/trtllm_lora_temp/data/csrc/fused_permute_quant.cuh create mode 100644 python/sglang/kernels/ops/moe/trtllm_lora_temp/data/csrc/trtllm_fused_moe_kernel_launcher.cu create mode 100644 python/sglang/kernels/ops/moe/trtllm_lora_temp/data/csrc/trtllm_fused_moe_runner.cu create mode 100644 python/sglang/kernels/ops/moe/trtllm_lora_temp/data/include/flashinfer/trtllm/fused_moe/DevKernel.h create mode 100644 python/sglang/kernels/ops/moe/trtllm_lora_temp/data/include/flashinfer/trtllm/fused_moe/runner.h rename python/sglang/{jit_kernel => kernels/ops/moe}/trtllm_lora_temp/jit.py (100%) rename python/sglang/{jit_kernel => kernels/ops/moe}/trtllm_lora_temp/kimi_k2_moe_fused_gate.py (100%) rename python/sglang/{jit_kernel => kernels/ops/moe}/trtllm_lora_temp/moe_lora_merged_align.py (100%) rename python/sglang/{jit_kernel => kernels/ops/moe}/trtllm_lora_temp/topk_pack.py (100%) rename python/sglang/{jit_kernel => kernels/ops/moe}/trtllm_lora_temp/topk_softmax_pack.py (100%) rename python/sglang/{jit_kernel => kernels/ops/quantization}/dsv32/__init__.py (100%) rename python/sglang/{jit_kernel => kernels/ops/quantization}/dsv32/elementwise.py (100%) diff --git a/python/sglang/jit_kernel/benchmark/kv_canary/utils.py b/python/sglang/jit_kernel/benchmark/kv_canary/utils.py index 80b76170a..7c74895d2 100644 --- a/python/sglang/jit_kernel/benchmark/kv_canary/utils.py +++ b/python/sglang/jit_kernel/benchmark/kv_canary/utils.py @@ -5,7 +5,7 @@ from typing import Callable import torch -from sglang.jit_kernel.kv_canary.verify import CANARY_SLOT_BYTES, RealKvSource +from sglang.kernels.ops.kv_canary.verify import CANARY_SLOT_BYTES, RealKvSource BS_AXIS: list[int] = [1, 4, 32, 128, 256, 1024] PREFIX_AXIS: list[int] = [0, 128, 1024, 4096, 10240, 16384] diff --git a/python/sglang/jit_kernel/dsa/__init__.py b/python/sglang/jit_kernel/dsa/__init__.py deleted file mode 100644 index 8509eb648..000000000 --- a/python/sglang/jit_kernel/dsa/__init__.py +++ /dev/null @@ -1,30 +0,0 @@ -from .paged_mqa_logits import ( - aiter_paged_mqa_logits, - cutedsl_paged_mqa_logits, - deepgemm_paged_mqa_logits_native, - deepgemm_paged_mqa_logits_split, -) - - -def pick_dsl_expand(*args, **kwargs): - from .cutedsl_paged_mqa_logits import pick_dsl_expand as _pick_dsl_expand - - return _pick_dsl_expand(*args, **kwargs) - - -def __getattr__(name: str): - if name == "CuteDSLPagedMQALogitsRunner": - from .cutedsl_paged_mqa_logits import CuteDSLPagedMQALogitsRunner - - return CuteDSLPagedMQALogitsRunner - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - - -__all__ = [ - "CuteDSLPagedMQALogitsRunner", - "pick_dsl_expand", - "aiter_paged_mqa_logits", - "cutedsl_paged_mqa_logits", - "deepgemm_paged_mqa_logits_native", - "deepgemm_paged_mqa_logits_split", -] diff --git a/python/sglang/jit_kernel/dsv4/__init__.py b/python/sglang/jit_kernel/dsv4/__init__.py deleted file mode 100644 index a5b4d48b2..000000000 --- a/python/sglang/jit_kernel/dsv4/__init__.py +++ /dev/null @@ -1,63 +0,0 @@ -from .attn import ( - fused_store_cache, - get_paged_mqa_logits_metadata, - triton_create_paged_compress_data, -) -from .c128_cleanup import clear_unaccepted_c128_draft_states -from .compress import ( - CompressorDecodePlan, - CompressorPrefillPlan, - compress_forward, - compress_norm_rope_store, -) -from .compress_old import fused_norm_rope_inplace -from .elementwise import ( - fused_k_norm_rope_flashmla, - fused_q_indexer_rope_first_quant, - fused_q_indexer_rope_hadamard_fp4_quant, - fused_q_indexer_rope_hadamard_quant, - fused_q_norm_rope, - fused_rope_inplace, -) -from .fp8_wo_a import sglang_per_token_group_quant_fp8_dsv4_wo_a -from .gemm import linear_bf16_fp32 -from .moe import ( - hash_topk, - mask_topk_ids, - mega_moe_pre_dispatch, - silu_and_mul_clamp, - silu_and_mul_contig_post_quant, - silu_and_mul_masked_post_quant, -) -from .topk import plan_topk_v2, topk_transform_512, topk_transform_512_v2 -from .utils import make_name - -__all__ = [ - "CompressorDecodePlan", - "CompressorPrefillPlan", - "compress_forward", - "compress_norm_rope_store", - "clear_unaccepted_c128_draft_states", - "fused_norm_rope_inplace", - "fused_store_cache", - "fused_rope_inplace", - "fused_q_norm_rope", - "fused_q_indexer_rope_first_quant", - "fused_q_indexer_rope_hadamard_fp4_quant", - "fused_q_indexer_rope_hadamard_quant", - "fused_k_norm_rope_flashmla", - "sglang_per_token_group_quant_fp8_dsv4_wo_a", - "make_name", - "linear_bf16_fp32", - "get_paged_mqa_logits_metadata", - "triton_create_paged_compress_data", - "topk_transform_512", - "topk_transform_512_v2", - "plan_topk_v2", - "hash_topk", - "mega_moe_pre_dispatch", - "mask_topk_ids", - "silu_and_mul_clamp", - "silu_and_mul_masked_post_quant", - "silu_and_mul_contig_post_quant", -] diff --git a/python/sglang/jit_kernel/kv_canary/plan/__init__.py b/python/sglang/jit_kernel/kv_canary/plan/__init__.py deleted file mode 100644 index e7bcc5c1d..000000000 --- a/python/sglang/jit_kernel/kv_canary/plan/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from sglang.jit_kernel.kv_canary.plan.api import launch_canary_plan_kernels diff --git a/python/sglang/jit_kernel/minimax_m3/__init__.py b/python/sglang/jit_kernel/minimax_m3/__init__.py deleted file mode 100644 index bbc930f95..000000000 --- a/python/sglang/jit_kernel/minimax_m3/__init__.py +++ /dev/null @@ -1,25 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Fused Triton kernels for MiniMax-M3 on AMD ROCm (gfx94x / gfx95x). - -Model-scoped JIT kernels (mirrors ``jit_kernel/dsv4``), split by op type: - * ``rmsnorm`` -- fused fp32 Gemma RMSNorm (plain + fused-add-residual) - * ``swiglu`` -- fused fp32 SwiGLU-OAI (split layout) -""" - -from sglang.jit_kernel.minimax_m3.rmsnorm import ( - _num_warps, - gemma_fused_add_rmsnorm, - gemma_rmsnorm, -) -from sglang.jit_kernel.minimax_m3.swiglu import ( - swiglu_oai_mxfp8_quant, - swiglu_oai_split, -) - -__all__ = [ - "gemma_rmsnorm", - "gemma_fused_add_rmsnorm", - "swiglu_oai_split", - "swiglu_oai_mxfp8_quant", - "_num_warps", -] diff --git a/python/sglang/jit_kernel/tests/deepseek_v4/common.py b/python/sglang/jit_kernel/tests/deepseek_v4/common.py index 66ad4014e..e16041ffc 100644 --- a/python/sglang/jit_kernel/tests/deepseek_v4/common.py +++ b/python/sglang/jit_kernel/tests/deepseek_v4/common.py @@ -5,7 +5,10 @@ from typing import List, Literal, Optional, Tuple import torch -from sglang.jit_kernel.dsv4 import CompressorDecodePlan, CompressorPrefillPlan +from sglang.kernels.ops.attention.dsv4 import ( + CompressorDecodePlan, + CompressorPrefillPlan, +) from sglang.srt.utils import get_device diff --git a/python/sglang/jit_kernel/tests/kv_canary/_canary_helpers.py b/python/sglang/jit_kernel/tests/kv_canary/_canary_helpers.py index 25b56f8c0..c7e67ace1 100644 --- a/python/sglang/jit_kernel/tests/kv_canary/_canary_helpers.py +++ b/python/sglang/jit_kernel/tests/kv_canary/_canary_helpers.py @@ -5,10 +5,6 @@ from typing import Optional import torch -from sglang.jit_kernel.kv_canary import consts -from sglang.jit_kernel.kv_canary.consts import splitmix64, splitmix64_mix3 -from sglang.jit_kernel.kv_canary.verify import VerifyPlan -from sglang.jit_kernel.kv_canary.write import WritePlan from sglang.jit_kernel.tests.kv_canary._constants import ( _I64_SIGN_BIT, _U64_MASK, @@ -20,6 +16,10 @@ from sglang.jit_kernel.tests.kv_canary._fixtures import ( make_real_kv_source, make_real_kv_sources, ) +from sglang.kernels.ops.kv_canary import consts +from sglang.kernels.ops.kv_canary.consts import splitmix64, splitmix64_mix3 +from sglang.kernels.ops.kv_canary.verify import VerifyPlan +from sglang.kernels.ops.kv_canary.write import WritePlan __all__ = [ "FakeViolationLog", diff --git a/python/sglang/jit_kernel/tests/kv_canary/_constants.py b/python/sglang/jit_kernel/tests/kv_canary/_constants.py index 5fb8874be..7313e9b89 100644 --- a/python/sglang/jit_kernel/tests/kv_canary/_constants.py +++ b/python/sglang/jit_kernel/tests/kv_canary/_constants.py @@ -1,6 +1,6 @@ from __future__ import annotations -from sglang.jit_kernel.kv_canary.verify import CANARY_SLOT_BYTES +from sglang.kernels.ops.kv_canary.verify import CANARY_SLOT_BYTES # Default fixture sizes — small enough for fast tests, large enough that ring overflow / multi-req cases # stay realistic without bloating the assertion surface. diff --git a/python/sglang/jit_kernel/tests/kv_canary/_differential.py b/python/sglang/jit_kernel/tests/kv_canary/_differential.py index 8c7c422fc..8a0ac1a7e 100644 --- a/python/sglang/jit_kernel/tests/kv_canary/_differential.py +++ b/python/sglang/jit_kernel/tests/kv_canary/_differential.py @@ -5,31 +5,31 @@ from typing import Any, Callable, Iterator, Optional import torch -from sglang.jit_kernel.kv_canary import consts -from sglang.jit_kernel.kv_canary.plan import launch_canary_plan_kernels -from sglang.jit_kernel.kv_canary.plan_ref import ( +from sglang.jit_kernel.tests.kv_canary._canary_helpers import ( + FakeViolationLog, + assert_canary_buf_equal, + assert_canary_state_equal, + make_log_pair, +) +from sglang.kernels.ops.kv_canary import consts +from sglang.kernels.ops.kv_canary.plan import launch_canary_plan_kernels +from sglang.kernels.ops.kv_canary.plan_ref import ( launch_canary_plan_kernels_torch_reference, ) -from sglang.jit_kernel.kv_canary.verify import ( +from sglang.kernels.ops.kv_canary.verify import ( CanaryLaunchTag, RealKvSource, VerifyOrWriteContext, VerifyPlan, launch_canary_verify_kernel, ) -from sglang.jit_kernel.kv_canary.verify_ref import ( +from sglang.kernels.ops.kv_canary.verify_ref import ( launch_canary_verify_kernel_torch_reference, ) -from sglang.jit_kernel.kv_canary.write import WritePlan, launch_canary_write_kernel -from sglang.jit_kernel.kv_canary.write_ref import ( +from sglang.kernels.ops.kv_canary.write import WritePlan, launch_canary_write_kernel +from sglang.kernels.ops.kv_canary.write_ref import ( launch_canary_write_kernel_torch_reference, ) -from sglang.jit_kernel.tests.kv_canary._canary_helpers import ( - FakeViolationLog, - assert_canary_buf_equal, - assert_canary_state_equal, - make_log_pair, -) _DEVICE = torch.device("cuda") diff --git a/python/sglang/jit_kernel/tests/kv_canary/_fixtures.py b/python/sglang/jit_kernel/tests/kv_canary/_fixtures.py index 98d6d21e1..2a04cb97a 100644 --- a/python/sglang/jit_kernel/tests/kv_canary/_fixtures.py +++ b/python/sglang/jit_kernel/tests/kv_canary/_fixtures.py @@ -5,12 +5,12 @@ from typing import Literal, Optional import torch -from sglang.jit_kernel.kv_canary.verify import ( +from sglang.jit_kernel.tests.kv_canary._constants import DEFAULT_NUM_SLOTS +from sglang.kernels.ops.kv_canary.verify import ( RealKvSource, VerifyPlan, ) -from sglang.jit_kernel.kv_canary.write import WritePlan -from sglang.jit_kernel.tests.kv_canary._constants import DEFAULT_NUM_SLOTS +from sglang.kernels.ops.kv_canary.write import WritePlan _DEVICE = torch.device("cuda") diff --git a/python/sglang/jit_kernel/tests/kv_canary/_hand_oracle.py b/python/sglang/jit_kernel/tests/kv_canary/_hand_oracle.py index 2f9487845..eebe7bf30 100644 --- a/python/sglang/jit_kernel/tests/kv_canary/_hand_oracle.py +++ b/python/sglang/jit_kernel/tests/kv_canary/_hand_oracle.py @@ -4,7 +4,7 @@ diff comparison.""" from __future__ import annotations -from sglang.jit_kernel.kv_canary.consts import splitmix64 +from sglang.kernels.ops.kv_canary.consts import splitmix64 def _fold_words(padded: bytes) -> int: diff --git a/python/sglang/jit_kernel/tests/kv_canary/_invariants.py b/python/sglang/jit_kernel/tests/kv_canary/_invariants.py index 55a7a6387..2bc7621c6 100644 --- a/python/sglang/jit_kernel/tests/kv_canary/_invariants.py +++ b/python/sglang/jit_kernel/tests/kv_canary/_invariants.py @@ -11,10 +11,10 @@ from typing import Optional import torch -from sglang.jit_kernel.kv_canary import consts -from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag, VerifyPlan -from sglang.jit_kernel.kv_canary.write import WritePlan from sglang.jit_kernel.tests.kv_canary._canary_helpers import FakeViolationLog +from sglang.kernels.ops.kv_canary import consts +from sglang.kernels.ops.kv_canary.verify import CanaryLaunchTag, VerifyPlan +from sglang.kernels.ops.kv_canary.write import WritePlan class PlanInvariants: diff --git a/python/sglang/jit_kernel/tests/test_minimax_m3_mxfp8.py b/python/sglang/jit_kernel/tests/test_minimax_m3_mxfp8.py index c8822e810..3fff5425e 100644 --- a/python/sglang/jit_kernel/tests/test_minimax_m3_mxfp8.py +++ b/python/sglang/jit_kernel/tests/test_minimax_m3_mxfp8.py @@ -86,7 +86,7 @@ def test_minimax_swiglu_mxfp8_quant_matches_unfused_fp32(m, inter): # the reference is the unfused fp32 swiglu followed by MXFP8 quant. Not # bit-identical because the reference quant runs in torch vs the fused triton # path, but numerically equivalent (tight relerr, scales agree within 1 ulp). - from sglang.jit_kernel.minimax_m3 import ( + from sglang.kernels.ops.moe.minimax_m3_swiglu import ( swiglu_oai_mxfp8_quant, swiglu_oai_split, ) diff --git a/python/sglang/jit_kernel/tests/test_minimax_m3_rmsnorm.py b/python/sglang/jit_kernel/tests/test_minimax_m3_rmsnorm.py index 6897039c8..7eb477fec 100644 --- a/python/sglang/jit_kernel/tests/test_minimax_m3_rmsnorm.py +++ b/python/sglang/jit_kernel/tests/test_minimax_m3_rmsnorm.py @@ -14,7 +14,7 @@ if not is_hip(): if not torch.cuda.is_available(): pytest.skip("Requires a GPU.", allow_module_level=True) -from sglang.jit_kernel.minimax_m3.rmsnorm import ( # noqa: E402 +from sglang.kernels.ops.layernorm.minimax_m3_rmsnorm import ( # noqa: E402 gemma_fused_add_rmsnorm, gemma_rmsnorm, ) diff --git a/python/sglang/jit_kernel/tests/test_moe_topk_sigmoid.py b/python/sglang/jit_kernel/tests/test_moe_topk_sigmoid.py index 4acbd23bd..769dacd92 100644 --- a/python/sglang/jit_kernel/tests/test_moe_topk_sigmoid.py +++ b/python/sglang/jit_kernel/tests/test_moe_topk_sigmoid.py @@ -13,7 +13,7 @@ from typing import Optional import pytest import torch -from sglang.jit_kernel.moe_topk_sigmoid import topk_sigmoid +from sglang.kernels.ops.moe.moe_topk_sigmoid import topk_sigmoid try: from sgl_kernel import topk_sigmoid as topk_sigmoid_aot diff --git a/python/sglang/jit_kernel/trtllm_lora_temp/__init__.py b/python/sglang/jit_kernel/trtllm_lora_temp/__init__.py deleted file mode 100644 index 5b2b4c1e2..000000000 --- a/python/sglang/jit_kernel/trtllm_lora_temp/__init__.py +++ /dev/null @@ -1,19 +0,0 @@ -from sglang.jit_kernel.trtllm_lora_temp.core import ( - trtllm_bf16_routed_moe_lora, - trtllm_fp4_block_scale_moe_lora_finalize, - trtllm_fp4_block_scale_routed_moe_lora, - trtllm_fp8_block_scale_moe, - trtllm_fp8_block_scale_moe_lora_finalize, - trtllm_fp8_block_scale_routed_moe, - trtllm_fp8_block_scale_routed_moe_lora, -) - -__all__ = [ - "trtllm_bf16_routed_moe_lora", - "trtllm_fp4_block_scale_moe_lora_finalize", - "trtllm_fp4_block_scale_routed_moe_lora", - "trtllm_fp8_block_scale_moe_lora_finalize", - "trtllm_fp8_block_scale_moe", - "trtllm_fp8_block_scale_routed_moe", - "trtllm_fp8_block_scale_routed_moe_lora", -] diff --git a/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_activation_quant.cuh b/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_activation_quant.cuh deleted file mode 100644 index f6b75710e..000000000 --- a/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_activation_quant.cuh +++ /dev/null @@ -1,246 +0,0 @@ -// Fused SwiGLU+LoRA activation -> NVFP4 per-token quant for the FP4 MoE LoRA path. -// -// Modeled on tensorrt_llm::kernels::nvfp4QuantAndPerTokenScaleKernel (flashinfer -// quantization.cuh): one block per expanded row, per-token amax via cub::BlockReduce, -// cvt_warp_fp16_to_fp4 for the e4m3 block scale + e2m1 packing + swizzled SF layout. -// The ONLY change vs that kernel is the pass-1 input: instead of reading activated_bf16 -// from gmem, it reads the interleaved gate/up GEMM1 output + the LoRA delta and computes -// silu(up)*gate on the fly, rounds to bf16 (matching the standalone activation kernel's -// bf16 output exactly), caches in smem, and also writes activation_lora_input. The down -// GEMM input (activated_bf16) is therefore never materialized to HBM. -// -// Because the activated value is rounded to bf16 before quantization (same as the separate -// activation kernel) and silu matches, the fp4 / SF / per_token_sf / activation_lora_input -// outputs are BITWISE-identical to the unfused activation -> quant#2 chain. -#pragma once - -#include -#include - -#include "nv_internal/tensorrt_llm/kernels/quantization_utils.cuh" -#include -#include -#include - -namespace flashinfer { -namespace sgl_fused_act_quant { - -namespace tk = tensorrt_llm::kernels; - -// Same silu as moe::dev::activation (trtllm_fused_moe_dev_kernel.cu:55): x / (1 + exp(-x)). -inline __device__ float fused_silu(float x) { - return x / (1.0f + expf(-x)); -} - -// One block per expanded row. gateUp is the column-interleaved GEMM1 output (g0,u0,g1,u1,...) -// indexed by permutedIdx; loraDelta is the contiguous [gate|up] delta indexed by expandedIdx. -template -__global__ void fusedActivationQuantKernel( - int m, // numTokens * topK (number of expanded rows) - int innerHalf, // inter == n (output width per row); must be a multiple of 16 - int innerDim, // gate_up_n == 2 * innerHalf - __nv_bfloat16 const* __restrict__ gateUp, // interleaved gate/up, [.., innerDim] by permutedIdx - __nv_bfloat16 const* __restrict__ loraDelta, // [.., innerDim] by expandedIdx, may be null - __nv_bfloat16* __restrict__ loraInputOut, // [.., innerHalf] by expandedIdx, may be null - int32_t const* __restrict__ expandedIdxToPermutedIdx, - float globalScaleInv, - uint8_t* __restrict__ weightOutput, // fp4 [.., innerHalf/2] by permutedIdx - uint8_t* __restrict__ scaleOutput, // swizzled e4m3 SF - float* __restrict__ perTokenScaleOutput) { - constexpr int SF_VEC_SIZE = 16; - using InType = tk::PackedVec<__nv_bfloat16, SF_VEC_SIZE>; // 16 bf16 == 8 __nv_bfloat162 - using PackedFp4Type = uint64_t; // SF_VEC_SIZE == 16 - - int const expandedIdx = blockIdx.x; - if (expandedIdx >= m) return; - int const permutedIdx = expandedIdxToPermutedIdx[expandedIdx]; - int const num_vecs_per_row = innerHalf / SF_VEC_SIZE; - int64_t const liBaseRow = (int64_t)expandedIdx * innerHalf; - - // Padding row: the separate activation kernel writes 0 to activation_lora_input and skips - // the quant outputs. Mirror that, then return. - if (permutedIdx < 0) { - if (loraInputOut != nullptr) { - InType z; -#pragma unroll - for (int i = 0; i < SF_VEC_SIZE / 2; ++i) - z.elts[i] = __float2bfloat162_rn(0.0f); - for (int vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) { - *reinterpret_cast(&loraInputOut[liBaseRow + (int64_t)vecIdx * SF_VEC_SIZE]) = z; - } - } - return; - } - - int64_t const permBase = (int64_t)permutedIdx * innerDim; // gate_up row (interleaved) - int64_t const expBase = (int64_t)expandedIdx * innerDim; // delta row (contiguous gate|up) - (void)DISABLE_FP4_FAST_MATH; - - // 1 SF block (16 outputs) per thread, held in registers across the amax barrier (no smem cache): - // requires num_vecs_per_row <= BLOCK_SIZE (inter=2048 -> 128 == BLOCK_SIZE). With - // CVT_ELTS_PER_THREAD == SF_VEC_SIZE the cvt needs no cross-thread shuffle, so masking is safe. - int const vecIdx = threadIdx.x; - bool const active = vecIdx < num_vecs_per_row; - - InType vec; - float localAmax = 0.f; - if (active) { - int const h0 = vecIdx * SF_VEC_SIZE; - __nv_bfloat16 const* g = gateUp + permBase + (int64_t)2 * h0; // 32 interleaved bf16 - __nv_bfloat16 const* dlo = loraDelta + expBase + h0; // silu-arg delta (lower half) - __nv_bfloat16 const* dhi = loraDelta + expBase + innerHalf + h0; // multiplier delta (upper half) - __nv_bfloat162 amax2 = __float2bfloat162_rn(0.0f); - union { - int4 v[4]; - __nv_bfloat16 b[32]; - } gu; - union { - int4 v[2]; - __nv_bfloat16 b[16]; - } dl, dh; - int4 const* gp = reinterpret_cast(g); -#pragma unroll - for (int k = 0; k < 4; ++k) - gu.v[k] = gp[k]; - if (loraDelta != nullptr) { - int4 const* dlp = reinterpret_cast(dlo); - int4 const* dhp = reinterpret_cast(dhi); -#pragma unroll - for (int k = 0; k < 2; ++k) { - dl.v[k] = dlp[k]; - dh.v[k] = dhp[k]; - } - } -#pragma unroll - for (int i = 0; i < SF_VEC_SIZE / 2; ++i) { // 8 bf162 = 16 output elements - int const j0 = 2 * i, j1 = 2 * i + 1; - float even0 = (float)gu.b[2 * j0], odd0 = (float)gu.b[2 * j0 + 1]; - float even1 = (float)gu.b[2 * j1], odd1 = (float)gu.b[2 * j1 + 1]; - float a0 = odd0, b0 = even0, a1 = odd1, b1 = even1; - if (loraDelta != nullptr) { - a0 += (float)dl.b[j0]; - b0 += (float)dh.b[j0]; - a1 += (float)dl.b[j1]; - b1 += (float)dh.b[j1]; - } - float act0 = fused_silu(a0) * b0; - float act1 = fused_silu(a1) * b1; - __nv_bfloat162 e = __float22bfloat162_rn(make_float2(act0, act1)); - vec.elts[i] = e; - amax2 = __hmax2(amax2, __habs2(e)); - } - localAmax = (float)__hmax(amax2.x, amax2.y); - if (loraInputOut != nullptr) { - *reinterpret_cast(&loraInputOut[liBaseRow + h0]) = vec; - } - } - - // ---- per-token scale: blockReduce amax, broadcast via smem (no gmem round-trip) ---- - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage tempStorage; - __shared__ float sScale; - float const globalAmax = BlockReduce(tempStorage).Reduce(localAmax, cuda::maximum<>{}); - if (threadIdx.x == 0) { - float const pts = globalAmax * globalScaleInv; - perTokenScaleOutput[permutedIdx] = pts; - sScale = pts; - } - __syncthreads(); - float const globalEncodeScale = tk::reciprocal_approximate_ftz(sScale); - - // ---- quantize from registers (cvt computes the per-16 e4m3 block scale internally) ---- - if (active) { - uint8_t fp8Scale; - // 5 template args on this flashinfer build: Type, SF_VEC_SIZE, CVT_ELTS_PER_THREAD, - // UE8M0_SF=false, TE_EXACT_NVFP4=false (the default nvfp4 quant path). - auto fp4Vals = tk::cvt_warp_fp16_to_fp4<__nv_bfloat16, SF_VEC_SIZE, SF_VEC_SIZE, false, false>( - vec, globalEncodeScale, &fp8Scale); - int64_t const vecOffset = (int64_t)permutedIdx * num_vecs_per_row + vecIdx; - reinterpret_cast(weightOutput)[vecOffset] = fp4Vals; - - // Match nvfp4QuantAndPerTokenScaleKernel exactly (it passes the kernel's `m` as numRows). - int64_t sfOffset; - if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::LINEAR) { - sfOffset = vecOffset; - } else if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) { - sfOffset = tk::get_sf_out_offset_128x4(/*batchIdx=*/0, permutedIdx, vecIdx, m, num_vecs_per_row); - } else { - sfOffset = tk::get_sf_out_offset_8x4(/*batchIdx=*/0, permutedIdx, vecIdx, m, num_vecs_per_row); - } - scaleOutput[sfOffset] = fp8Scale; - } -} - -// Host launch: globalScaleInv = 1/448/6. BLOCK_SIZE must be >= innerHalf/16 (one SF block/thread). -inline void launchFusedActivationQuant( - int m, - int innerHalf, - int innerDim, - __nv_bfloat16 const* gateUp, - __nv_bfloat16 const* loraDelta, - __nv_bfloat16* loraInputOut, - int32_t const* expandedIdxToPermutedIdx, - float globalScaleInv, - uint8_t* weightOutput, - uint8_t* scaleOutput, - float* perTokenScaleOutput, - tensorrt_llm::QuantizationSFLayout sfLayout, - bool disableFp4FastMath, - cudaStream_t stream) { - // One SF block per thread, no stride loop: BLOCK_SIZE must cover innerHalf/16 (a fixed - // 128 left cols [2048,inter) unwritten at Inkling EP8's inter=3072 -> NaN from the down GEMM). - uint32_t const numVecs = static_cast(innerHalf) / 16; - auto dispatchBlock = [&](auto blockTag) { - constexpr uint32_t BLOCK_SIZE = decltype(blockTag)::value; - dim3 const grid(m), block(BLOCK_SIZE); - auto launch = [&](auto layoutTag, auto fastMathTag) { - fusedActivationQuantKernel - <<>>( - m, - innerHalf, - innerDim, - gateUp, - loraDelta, - loraInputOut, - expandedIdxToPermutedIdx, - globalScaleInv, - weightOutput, - scaleOutput, - perTokenScaleOutput); - }; - auto withFastMath = [&](auto layoutTag) { - if (disableFp4FastMath) { - launch(layoutTag, std::integral_constant{}); - } else { - launch(layoutTag, std::integral_constant{}); - } - }; - if (sfLayout == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) { - withFastMath( - std::integral_constant< - tensorrt_llm::QuantizationSFLayout, - tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4>{}); - } else if (sfLayout == tensorrt_llm::QuantizationSFLayout::LINEAR) { - withFastMath( - std::integral_constant{}); - } else { - withFastMath( - std::integral_constant< - tensorrt_llm::QuantizationSFLayout, - tensorrt_llm::QuantizationSFLayout::SWIZZLED_8x4>{}); - } - }; - if (numVecs <= 128) { - dispatchBlock(std::integral_constant{}); - } else if (numVecs <= 256) { - dispatchBlock(std::integral_constant{}); - } else if (numVecs <= 512) { - dispatchBlock(std::integral_constant{}); - } else { - // Callers guard on numVecs <= 512 and fall back to the unfused chain. - dispatchBlock(std::integral_constant{}); - } -} - -} // namespace sgl_fused_act_quant -} // namespace flashinfer diff --git a/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_permute_quant.cuh b/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_permute_quant.cuh deleted file mode 100644 index 4fed01297..000000000 --- a/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_permute_quant.cuh +++ /dev/null @@ -1,307 +0,0 @@ -// Fused permute + NvFP4-per-token-quant for the FP4 MoE-LoRA gate_up path. -// -// Background (decode bs64, EP8): the plain path runs `permuteKernel` (gather bf16 hidden into the -// padded [max_padded, hidden] permuted buffer) then `nvfp4QuantAndPerTokenScaleKernel` over ALL -// max_padded rows. At decode only num_tokens*top_k of those rows are real (the rest are padding), -// so both kernels waste ~6x of their work on padding, and the bf16 permuted buffer is a full -// HBM round-trip (written by permute, read back by quant). -// -// This kernel fuses the two: it reads the UN-permuted hidden, NvFP4-quantizes each (token,expert) -// pair's row, and scatter-writes fp4 + swizzled block-sf + per-token-sf directly to that pair's -// permuted position. It iterates only the num_tokens*top_k real pairs (skipping pad), and never -// materializes the bf16 permuted buffer. -// -// It mirrors `nvfp4QuantAndPerTokenScaleKernel` (quantization.cuh) — same amax, same per-token-scale -// recipe, same `cvt_warp_fp16_to_fp4`, same swizzled-sf offset (`get_sf_out_offset_8x4`) — with the -// single `rowIdx` split into a READ row (the unpermuted source token) and a WRITE row (the permuted -// destination). For the valid rows the result is BITWISE-identical to the plain permute->quant chain -// (the chain's quant reads permuted_hidden[writeRow], filled by permute from hidden[readRow]; we -// read hidden[readRow] directly), verified by the bench's fused-vs-old guard. -// -// PER-TOKEN-SCALE BRANCH: uses TE_EXACT (globalEncodeScale = __fdiv_rn(globalScale, globalAmax), -// stored scale = 1/globalEncodeScale, cvt TE_EXACT_NVFP4=true). This matches the installed -// flashinfer 0.6.11.post1, whose DISPATCH macro hard-codes the bf16 kernel to TE_EXACT_NVFP4=true -// (quantization.cu DISPATCH_NVP4_QUANT_AND_PER_TOKEN_SCALE_KERNEL). NOTE: if a future flashinfer -// reverts the bf16 path to the fast-math branch (reciprocal_approximate / TE_EXACT=false), or if -// FLASHINFER_NVFP4_4OVER6 is enabled, this fused kernel would diverge — re-validate the bench's -// fused-vs-old bitwise guard against the deployed flashinfer before trusting it there. -// -// Two variants (both kept, selectable, for cross-scenario perf comparison): -// - no-dedup: grid over the num_tokens*top_k pairs; each block re-reads+re-quantizes its source -// token and writes 1 destination (more blocks -> better occupancy at tiny decode sizes). -// - dedup: grid over num_tokens; each block reads+quantizes its token once and scatter-writes -// to all of that token's (valid) permuted destinations (no redundant quant, fewer blocks). -// -// Helpers are pulled from quantization_utils.cuh (cvt_warp_fp16_to_fp4 / get_sf_out_offset_* / -// PackedVec / reciprocal_approximate_ftz) rather than quantization.cuh, because the latter pulls in -// nv_internal/.../common/cudaUtils.h, which ODR-conflicts with the flashinfer/trtllm/common twin -// already in trtllm_fused_moe_kernel_launcher.cu's TU. loadPackedVec lives in quantization.cuh, so -// we do a direct aligned PackedVec load instead. -#pragma once - -#include - -#include "nv_internal/tensorrt_llm/kernels/quantization_utils.cuh" -#include -#include -#include -#include -#include - -namespace sgl_fused_permute_quant { - -namespace tk = tensorrt_llm::kernels; - -// Quantize source row `readRow` of `input` (unpermuted) and write fp4 + block-sf + per-token-sf to -// destination row `writeRow` of the permuted outputs. `numRowsSf` is the SF buffer's row count -// (= max_padded), matching the plain quant's `m` arg to get_sf_out_offset_*. -template -__device__ __forceinline__ void fused_quant_one_row( - uint32_t n, - T const* input, - int readRow, - int writeRow, - int numRowsSf, - float globalScaleInv, - uint8_t* weightOutput, - uint8_t* scaleOutput, - float* perTokenScaleOutput) { - constexpr int SF_VEC_SIZE = 16; - constexpr int ELTS_PER_THREAD = 16; - using InType = tk::PackedVec; - using PackedFp4Type = std::conditional_t; - uint32_t const num_vecs_per_row = (n + ELTS_PER_THREAD - 1) / ELTS_PER_THREAD; - uint32_t const num_sf_vecs_per_row = (n + SF_VEC_SIZE - 1) / SF_VEC_SIZE; - InType const* inBase = reinterpret_cast(input); - - // ---- pass 1: per-row amax over the (unpermuted) source row ---- - float localAmax = 0.f; - for (uint32_t vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) { - InType vec_in = inBase[static_cast(readRow) * num_vecs_per_row + vecIdx]; - std::remove_reference_t a(0.f, 0.f); -#pragma unroll - for (int i = 0; i < ELTS_PER_THREAD / 2; ++i) { - a = __hmax2(a, __habs2(vec_in.elts[i])); - } - localAmax = fmaxf(localAmax, static_cast(__hmax(a.x, a.y))); - } - - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage tempStorage; - float const globalAmax = BlockReduce(tempStorage).Reduce(localAmax, cuda::maximum<>{}); - - // ---- per-token scale (TE_EXACT branch — production instantiates TE_EXACT_NVFP4=true for bf16, - // quantization.cu:247): globalEncodeScale = globalScale/globalAmax (exact __fdiv_rn), stored - // per-token scale = 1/globalEncodeScale. __shared__ scalar replaces the gmem round-trip - // (bit-identical: an fp32 store->load doesn't change the value). ---- - __shared__ float sEncodeScale; - if (threadIdx.x == 0) { - float const globalScale = __fdiv_rn(1.0f, globalScaleInv); - float const rowEncodeScale = globalAmax != 0.0f ? fminf(__fdiv_rn(globalScale, globalAmax), FLT_MAX) : FLT_MAX; - sEncodeScale = rowEncodeScale != 0.0f ? rowEncodeScale : 1.0f; - } - __syncthreads(); - float const globalEncodeScale = sEncodeScale; - float const perTokenScale = __fdiv_rn(1.0f, globalEncodeScale); - if (threadIdx.x == 0) perTokenScaleOutput[writeRow] = perTokenScale; - - // ---- pass 2: quantize + scatter-write to the permuted destination ---- - for (uint32_t vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) { - InType vec_in = inBase[static_cast(readRow) * num_vecs_per_row + vecIdx]; - uint8_t fp8Scale; - auto fp4Vals = tk::cvt_warp_fp16_to_fp4< - T, - SF_VEC_SIZE, - ELTS_PER_THREAD, - /*UE8M0_SF=*/false, - /*TE_EXACT_NVFP4=*/true>(vec_in, globalEncodeScale, &fp8Scale); - reinterpret_cast(weightOutput)[static_cast(writeRow) * num_vecs_per_row + vecIdx] = - fp4Vals; - - int64_t sfOffset; - if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::LINEAR) { - sfOffset = static_cast(writeRow) * num_sf_vecs_per_row + vecIdx; - } else if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) { - sfOffset = tk::get_sf_out_offset_128x4(/*batchIdx=*/0, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row); - } else { - sfOffset = tk::get_sf_out_offset_8x4(/*batchIdx=*/0, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row); - } - scaleOutput[sfOffset] = fp8Scale; - } -} - -// no-dedup: grid.x = num_tokens*top_k (one block per (token,expert) pair). -template -__global__ void fusedPermuteNvfp4QuantKernel( - uint32_t numPairs, - uint32_t n, - uint32_t topK, - int numRowsSf, - T const* input, - float globalScaleInv, - int32_t const* expandedIdxToPermutedIdx, - uint8_t* weightOutput, - uint8_t* scaleOutput, - float* perTokenScaleOutput) { - uint32_t const expandedIdx = blockIdx.x; - if (expandedIdx >= numPairs) return; - int const writeRow = expandedIdxToPermutedIdx[expandedIdx]; - if (writeRow < 0) return; - int const readRow = static_cast(expandedIdx / topK); - fused_quant_one_row( - n, input, readRow, writeRow, numRowsSf, globalScaleInv, weightOutput, scaleOutput, perTokenScaleOutput); -} - -// dedup: grid.x = num_tokens (one block per source token, scatter to its top_k destinations). -template -__global__ void fusedPermuteNvfp4QuantDedupKernel( - uint32_t numTokens, - uint32_t n, - uint32_t topK, - int numRowsSf, - T const* input, - float globalScaleInv, - int32_t const* expandedIdxToPermutedIdx, - uint8_t* weightOutput, - uint8_t* scaleOutput, - float* perTokenScaleOutput) { - constexpr int SF_VEC_SIZE = 16; - constexpr int ELTS_PER_THREAD = 16; - using InType = tk::PackedVec; - using PackedFp4Type = std::conditional_t; - uint32_t const token = blockIdx.x; - if (token >= numTokens) return; - uint32_t const num_vecs_per_row = (n + ELTS_PER_THREAD - 1) / ELTS_PER_THREAD; - uint32_t const num_sf_vecs_per_row = (n + SF_VEC_SIZE - 1) / SF_VEC_SIZE; - InType const* inBase = reinterpret_cast(input); - - // pass 1: amax over the source token row (read once). - float localAmax = 0.f; - for (uint32_t vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) { - InType vec_in = inBase[static_cast(token) * num_vecs_per_row + vecIdx]; - std::remove_reference_t a(0.f, 0.f); -#pragma unroll - for (int i = 0; i < ELTS_PER_THREAD / 2; ++i) { - a = __hmax2(a, __habs2(vec_in.elts[i])); - } - localAmax = fmaxf(localAmax, static_cast(__hmax(a.x, a.y))); - } - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage tempStorage; - float const globalAmax = BlockReduce(tempStorage).Reduce(localAmax, cuda::maximum<>{}); - - // TE_EXACT per-token scale (matches production; see fused_quant_one_row). - __shared__ float sEncodeScale; - if (threadIdx.x == 0) { - float const globalScale = __fdiv_rn(1.0f, globalScaleInv); - float const rowEncodeScale = globalAmax != 0.0f ? fminf(__fdiv_rn(globalScale, globalAmax), FLT_MAX) : FLT_MAX; - sEncodeScale = rowEncodeScale != 0.0f ? rowEncodeScale : 1.0f; - } - __syncthreads(); - float const globalEncodeScale = sEncodeScale; - float const perTokenScale = __fdiv_rn(1.0f, globalEncodeScale); - - // per-token scale -> each (valid) destination (top_k small; first top_k threads write). - if (threadIdx.x < topK) { - int const writeRow = expandedIdxToPermutedIdx[token * topK + threadIdx.x]; - if (writeRow >= 0) perTokenScaleOutput[writeRow] = perTokenScale; - } - - // pass 2: quantize each vec once, scatter to all valid destinations. - for (uint32_t vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) { - InType vec_in = inBase[static_cast(token) * num_vecs_per_row + vecIdx]; - uint8_t fp8Scale; - auto fp4Vals = tk::cvt_warp_fp16_to_fp4< - T, - SF_VEC_SIZE, - ELTS_PER_THREAD, - /*UE8M0_SF=*/false, - /*TE_EXACT_NVFP4=*/true>(vec_in, globalEncodeScale, &fp8Scale); -#pragma unroll 1 - for (uint32_t k = 0; k < topK; ++k) { - int const writeRow = expandedIdxToPermutedIdx[token * topK + k]; - if (writeRow < 0) continue; - reinterpret_cast(weightOutput)[static_cast(writeRow) * num_vecs_per_row + vecIdx] = - fp4Vals; - int64_t sfOffset; - if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::LINEAR) { - sfOffset = static_cast(writeRow) * num_sf_vecs_per_row + vecIdx; - } else if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) { - sfOffset = tk::get_sf_out_offset_128x4(/*batchIdx=*/0, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row); - } else { - sfOffset = tk::get_sf_out_offset_8x4(/*batchIdx=*/0, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row); - } - scaleOutput[sfOffset] = fp8Scale; - } - } -} - -// Launcher. `dedup` picks the variant. `n` (= hidden) must be a multiple of 16. `numRowsSf` is the -// SF buffer's row count (max_padded). -template -void invokeFusedPermuteNvfp4Quant( - uint32_t numTokens, - uint32_t topK, - uint32_t n, - int numRowsSf, - T const* input, - float globalScaleInv, - int32_t const* expandedIdxToPermutedIdx, - uint8_t* weightOutput, - uint8_t* scaleOutput, - float* perTokenScaleOutput, - tensorrt_llm::QuantizationSFLayout sfLayout, - bool dedup, - cudaStream_t stream) { - // [opt] Occupancy tuning (ncu: kernel is occupancy-bound, not DRAM-bound — DRAM <1%, achieved - // occupancy was 19.8% no-dedup / 5.5% dedup at BLOCK_SIZE=128). The dedup variant launches only - // num_tokens CTAs (=64 at decode bs64), so it is the most CTA-starved; widening the block raises - // threads/CTA and hides the per-row amax-reduction + scatter latency. Decode bs64 dedup sweep: - // 128 -> 5.52us, 256 -> 4.09us, 512 -> 3.71us. 512 is the chosen default (the prod path uses - // dedup). (7168/16 = 448 vecs/row, so >448 threads idle on the tail, but the win dominates.) - constexpr uint32_t BLOCK_SIZE = 512; - dim3 const block(BLOCK_SIZE); - - auto dispatch = [&](auto layoutTag) { - constexpr tensorrt_llm::QuantizationSFLayout LAYOUT = decltype(layoutTag)::value; - if (dedup) { - dim3 const grid(numTokens); - fusedPermuteNvfp4QuantDedupKernel<<>>( - numTokens, - n, - topK, - numRowsSf, - input, - globalScaleInv, - expandedIdxToPermutedIdx, - weightOutput, - scaleOutput, - perTokenScaleOutput); - } else { - dim3 const grid(numTokens * topK); - fusedPermuteNvfp4QuantKernel<<>>( - numTokens * topK, - n, - topK, - numRowsSf, - input, - globalScaleInv, - expandedIdxToPermutedIdx, - weightOutput, - scaleOutput, - perTokenScaleOutput); - } - }; - - if (sfLayout == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) { - dispatch( - std::integral_constant< - tensorrt_llm::QuantizationSFLayout, - tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4>{}); - } else { - dispatch( - std::integral_constant{}); - } -} - -} // namespace sgl_fused_permute_quant diff --git a/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/trtllm_fused_moe_kernel_launcher.cu b/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/trtllm_fused_moe_kernel_launcher.cu deleted file mode 100644 index 8ed7c7ac9..000000000 --- a/python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/trtllm_fused_moe_kernel_launcher.cu +++ /dev/null @@ -1,4380 +0,0 @@ -/* - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT 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 - -#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmGatedActOptions.h" -#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h" -#include "flashinfer/trtllm/fused_moe/DevKernel.h" -#include "flashinfer/trtllm/fused_moe/RoutingKernel.h" -#include "flashinfer/trtllm/fused_moe/runner.h" -#include "fused_activation_quant.cuh" -#include "fused_permute_quant.cuh" // fused permute+nvfp4-quant (gate_up de-pad), used by bench_fused_permute_quant -#include "nv_internal/tensorrt_llm/kernels/quantization.h" -#include "nv_internal/tensorrt_llm/thop/utils.h" -#include "tvm_ffi_utils.h" -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace flashinfer { - -namespace btg = batchedGemm::trtllm::gen; -using tensorrt_llm::kernels::trtllmgen_moe::MoE::ActivationType; -using tensorrt_llm::kernels::trtllmgen_moe::Routing::RoutingMethodType; -using tvm::ffi::Array; -using tvm::ffi::Optional; - -enum class RoutingInputMode { - FromLogits, // Mode 1: Compute routing from logits - PackedPrecomputed, // Mode 2: Pre-computed with packed (score << 16 | id) format - UnpackedPrecomputed // Mode 3: Pre-computed with separate topk_ids and topk_weights -}; - -// Validate routing_replay_out tensor properties. -// NOTE: dim0 >= num_tokens is intentionally NOT checked — with CUDA graphs the buffer -// is pre-allocated at maximum batch size and reused across steps with varying num_tokens. -static void validate_routing_replay_out(TensorView const& replay, TensorView const& hidden_states, int64_t top_k) { - TVM_FFI_ICHECK(replay.device().device_type == kDLCUDA) << "routing_replay_out must be a CUDA tensor"; - TVM_FFI_ICHECK(replay.device().device_id == hidden_states.device().device_id) - << "routing_replay_out must be on the same device as hidden_states"; - TVM_FFI_ICHECK(replay.ndim() == 2) << "routing_replay_out must be 2D [num_tokens, top_k]"; - TVM_FFI_ICHECK(replay.size(1) == top_k) << "routing_replay_out dim1 must equal top_k"; - TVM_FFI_ICHECK((replay.dtype() == DLDataType{kDLInt, 16, 1})) << "routing_replay_out must be int16 dtype"; - TVM_FFI_ICHECK(replay.IsContiguous()) << "routing_replay_out must be contiguous (packed row-major)"; -} - -enum class Fp8QuantizationType { - NoneFp8, - DeepSeekFp8, - MxFp8, - PerTensorFp8, -}; - -inline std::string fp8QuantizationTypeToString(Fp8QuantizationType quantization_type) { - switch (quantization_type) { - default: - case Fp8QuantizationType::NoneFp8: - return "NoneFp8"; - case Fp8QuantizationType::DeepSeekFp8: - return "DeepSeekFp8"; - case Fp8QuantizationType::MxFp8: - return "MxFp8"; - case Fp8QuantizationType::PerTensorFp8: - return "PerTensorFp8"; - } -} - -inline ActivationType validateAndCastActivationType(int64_t act_type) { - TVM_FFI_ICHECK(act_type >= 0 && act_type < static_cast(ActivationType::InvalidType)) - << "Invalid activation type: " << act_type; - return static_cast(act_type); -} - -// Utility function to compute the next power of two -inline int32_t nextPowerOfTwo(float value) { - int32_t n = static_cast(std::ceil(value)); - if (n <= 1) return 1; - - // If n is already a power of 2, return it - if ((n & (n - 1)) == 0) return n; - - // Find the next power of 2 - n--; - n |= n >> 1; - n |= n >> 2; - n |= n >> 4; - n |= n >> 8; - n |= n >> 16; - n++; - - return n; -} - -std::set computeSelectedTileN( - std::vector const& supported_tile_nums, - int64_t const num_tokens, - int64_t const top_k, - int64_t const num_local_experts) { - TVM_FFI_ICHECK(!supported_tile_nums.empty()) << "supported_tile_nums must not be empty."; - float const avg_tokens_per_expert = static_cast(num_tokens * top_k) / num_local_experts; - // NOTE: This differs from Python AutoTuner bucketing: - // - AutoTuner maps raw num_tokens with last_positive_power_of_2 (round-down). - // - Here we map derived avg_tokens_per_expert and use nextPowerOfTwo (round-up). - // Because they round different quantities in different directions, cache bucket and runtime - // tile candidates can diverge; launcher-side tactic resolution handles that mismatch. - // assume supported_tile_nums is sorted - int32_t tile_tokens_dim = - std::clamp(nextPowerOfTwo(avg_tokens_per_expert), supported_tile_nums.front(), supported_tile_nums.back()); - auto it = std::find(supported_tile_nums.begin(), supported_tile_nums.end(), tile_tokens_dim); - FLASHINFER_CHECK( - it != supported_tile_nums.end(), - "computeSelectedTileN expected exact tile ", - tile_tokens_dim, - " in supported_tile_nums (size=", - supported_tile_nums.size(), - "). Please keep supported_tile_nums as a dense power-of-2 ladder for this launcher."); - - // Candidate tile set centered on the heuristic tile. - // This function returns nearby candidates (not a single final tile): - // center, +1, +2, and -1 neighbors when available. - // Final tile choice is made later (autotuner-provided tile if valid, otherwise fallback policy). - std::set selected_tile_nums; - selected_tile_nums.insert(tile_tokens_dim); - if (std::next(it) != supported_tile_nums.end()) { - selected_tile_nums.insert(*std::next(it)); - if (std::next(std::next(it)) != supported_tile_nums.end()) { - selected_tile_nums.insert(*std::next(std::next(it))); - } - } - if (it != supported_tile_nums.begin()) { - selected_tile_nums.insert(*std::prev(it)); - } - - return selected_tile_nums; -} - -int64_t selectDefaultTileN( - std::vector const& supported_tile_nums, - int64_t const num_tokens, - int64_t const top_k, - int64_t const num_local_experts) { - auto selected = computeSelectedTileN(supported_tile_nums, num_tokens, top_k, num_local_experts); - TVM_FFI_ICHECK(!selected.empty()) << "No selected tile_N candidates for current MoE input."; - return *selected.begin(); -} - -// Resolve the (tile_N, config) pair passed from Python side, applying fallback logic -// when tile_N is -1. -std::pair resolveMoeTileAndConfig( - Array const& config_index, - std::vector const& supported_tile_nums, - int64_t const num_tokens, - int64_t const top_k, - int64_t const num_local_experts) { - // Python side convention: tactic is [tile_N, config] - TVM_FFI_ICHECK(config_index.size() == 2) - << "Invalid tactic, expected to be [tile_N, config], but got array of size " << config_index.size(); - const int64_t tile_N = config_index[0]; - const int64_t config = config_index[1]; - - if (tile_N == -1 || config == -1) { - // Use fallback tactic - auto const default_tile_N = selectDefaultTileN(supported_tile_nums, num_tokens, top_k, num_local_experts); - return {default_tile_N, -1}; - } - - return {tile_N, config}; -} - -class FusedMoeLauncher { - protected: - Optional routing_logits; - Optional routing_bias; - TensorView hidden_states; - TensorView gemm1_weights; - Optional output1_scales_scalar; - Optional output1_scales_gate_scalar; - TensorView gemm2_weights; - Optional output2_scales_scalar; - Optional per_token_scales; - Tensor per_token_scales_fc2; - - int64_t tile_tokens_dim{}; - int64_t routing_method_type{}; - bool use_shuffled_weight{}; - batchedGemm::gemm::MatrixLayout weight_layout{batchedGemm::gemm::MatrixLayout::MajorK}; - - std::tuple device_version; - std::unique_ptr args; - tensorrt_llm::kernels::trtllmgen_moe::MoE::MoEWorkspace workspace; - - btg::Dtype mDtypeAct{btg::Dtype::Bfloat16}; - btg::Dtype mDtypeWeights{btg::Dtype::Bfloat16}; - btg::Dtype mRoutingBiasDtype{btg::Dtype::Bfloat16}; // Dtype for expert weights in routing, based on routing bias - btg::Dtype mRoutingLogitsDtype{btg::Dtype::Bfloat16}; - bool norm_topk_prob{true}; - ActivationType activation_type{ActivationType::Swiglu}; - btg::Dtype mDtypeScore{btg::Dtype::Bfloat16}; - - // Optional routing replay output: [num_tokens, top_k] int16 tensor - Optional routing_replay_out; - - int64_t intermediate_size_factor{2}; - - public: - // Constructor that initializes all TensorView members - FusedMoeLauncher( - const Optional& routing_logits, - const Optional& routing_bias, - const TensorView& hidden_states, - const TensorView& gemm1_weights, - const Optional& output1_scales_scalar, - const Optional& output1_scales_gate_scalar, - const TensorView& gemm2_weights, - const Optional& output2_scales_scalar, - const Optional& per_token_scales) - : routing_logits(routing_logits), - routing_bias(routing_bias), - hidden_states(hidden_states), - gemm1_weights(gemm1_weights), - output1_scales_scalar(output1_scales_scalar), - output1_scales_gate_scalar(output1_scales_gate_scalar), - gemm2_weights(gemm2_weights), - output2_scales_scalar(output2_scales_scalar), - per_token_scales(per_token_scales), - tile_tokens_dim{}, - routing_method_type{}, - use_shuffled_weight{}, - weight_layout{batchedGemm::gemm::MatrixLayout::MajorK}, - mDtypeAct{btg::Dtype::Bfloat16}, - mDtypeWeights{btg::Dtype::Bfloat16}, - activation_type{ActivationType::Swiglu}, - intermediate_size_factor{2} {} - - public: - void set_routing_replay_out(const Optional& replay_out) { - routing_replay_out = replay_out; - } - - protected: - // Initialize common data necessary for later. - // May throw exception from TVM_FFI_ICHECK. - void init_common( - std::unique_ptr&& args, - int64_t tile_tokens_dim, - int64_t routing_method_type, - bool use_shuffled_weight, - int64_t weight_layout, - ActivationType activation_type, - bool norm_topk_prob = true); - - // Routing logits [num_tokens, num_experts] - void check_routing_logits() const { - if (routing_logits.has_value()) { - // Check shape - TVM_FFI_ICHECK_EQ(routing_logits.value().ndim(), 2) << "routing_logits must be 2D."; - TVM_FFI_ICHECK_EQ(routing_logits.value().size(0), hidden_states.size(0)) - << "routing_logits and hidden_states must have the same number of tokens."; - TVM_FFI_ICHECK_EQ(routing_logits.value().size(1), args->num_experts) - << "routing_logits dim1 must match num_experts."; - - // Check dtype - TVM_FFI_ICHECK(routing_logits.value().dtype() == dl_float32 || routing_logits.value().dtype() == dl_bfloat16) - << "routing_logits must be float or bfloat16."; - } - } - - // Routing bias [num_experts] - void check_routing_bias_shape() const { - if (routing_bias.has_value()) { - TVM_FFI_ICHECK_EQ(routing_bias.value().ndim(), 1) << "routing_bias must be 1D."; - TVM_FFI_ICHECK_EQ(routing_bias.value().size(0), args->num_experts) << "routing_bias has incorrect shape."; - } - } - - // Hidden states [num_tokens, hidden_size] - void check_hidden_states_shape() const { - TVM_FFI_ICHECK_EQ(hidden_states.ndim(), 2) << "hidden_states must be 2D."; - TVM_FFI_ICHECK_EQ(hidden_states.size(1), args->intermediate_size) << "hidden_states has incorrect shape."; - } - - // GEMM1 or GEMM2 weights [num_experts, M, K] or [num_experts, K/block_k, M, block_k] - void check_weights_shape(std::string which_weights) const { - TensorView weights = (which_weights == "gemm1") ? gemm1_weights : gemm2_weights; - if (which_weights != "gemm1" && which_weights != "gemm2") { - TVM_FFI_LOG_AND_THROW(InternalError) << "Internal error: which_weights = " << which_weights; - } - - int64_t Mn = 0, K = 0; - if (weight_layout == batchedGemm::gemm::MatrixLayout::MajorK) { - // MajorK [num_experts, M, K] - Mn = weights.size(1); - K = weights.size(2); - } else if (weight_layout == batchedGemm::gemm::MatrixLayout::BlockMajorK) { - // BlockMajorK [num_experts, K/block_k, M, block_k] - Mn = weights.size(2); - int64_t block_k = weights.size(3); - K = weights.size(1) * block_k; - } else { - TVM_FFI_LOG_AND_THROW(NotImplementedError) << "Unsupported weight_layout: " << (int)weight_layout; - } - if (which_weights == "gemm1") { - // Gated MoE activations (e.g. Swiglu/Geglu) pack gate+up projections in GEMM1, - // so Mn = 2 * intermediate_size and must be even. - if (intermediate_size_factor == 2) { - TVM_FFI_ICHECK_EQ(Mn % 2, 0) << which_weights << " weights Mn dimension must be even."; - } - // Non-gated activations (e.g. Relu2) use a single projection in GEMM1, - // so Mn = intermediate_size. This check covers both gated and non-gated cases. - TVM_FFI_ICHECK_EQ(args->intermediate_size * intermediate_size_factor, Mn) - << "intermediate_size has incorrect shape."; - TVM_FFI_ICHECK_EQ(K, hidden_states.size(1)) - << which_weights << " weights K dimension must be equal to hidden_size."; - } else if (which_weights == "gemm2") { - // GEMM2 always consumes the post-activation hidden of size intermediate_size. - TVM_FFI_ICHECK_EQ(K, args->intermediate_size) - << which_weights << " weights K dimension must be equal to intermediate_size."; - } - } - - void check_routing_common() const { - TVM_FFI_ICHECK(args->top_k > 0 && args->top_k <= args->num_experts) << "top_k must be between 1 and num_experts"; - TVM_FFI_ICHECK(args->local_num_experts > 0 && args->local_num_experts <= args->num_experts) - << "local_num_experts must be between 1 and num_experts"; - TVM_FFI_ICHECK( - args->local_expert_offset >= 0 && args->local_expert_offset + args->local_num_experts <= args->num_experts) - << "expert offset and count must be within valid range"; - - check_routing_logits(); - - if (routing_bias.has_value()) { - check_routing_bias_shape(); - } - } - - // Routing phase workspace tensors (allocated in prepare_routing() or prepare_routing_common()) - Tensor num_tokens_per_expert; - Tensor total_num_padded_tokens; - Tensor expanded_idx_to_permuted_idx; - Tensor permuted_idx_to_token_idx; - Tensor expert_weights; - Tensor expert_indexes; - Tensor expert_count_histogram; - Tensor cta_idx_xy_to_batch_idx; - Tensor cta_idx_xy_to_mn_limit; - Tensor num_non_exiting_ctas; - - void prepare_routing_common() { - // Allocate routing phase workspace tensors - num_tokens_per_expert = alloc_tensor({args->num_experts}, dl_int32, hidden_states.device()); - int32_t max_num_padded_tokens = tensorrt_llm::kernels::trtllmgen_moe::Routing::getMaxPermutedPaddedCount( - args->num_tokens, args->top_k, args->num_experts, tile_tokens_dim); - - total_num_padded_tokens = alloc_tensor({1}, dl_int32, hidden_states.device()); - - expanded_idx_to_permuted_idx = alloc_tensor({args->num_tokens * args->top_k}, dl_int32, hidden_states.device()); - - permuted_idx_to_token_idx = alloc_tensor({max_num_padded_tokens}, dl_int32, hidden_states.device()); - - expert_indexes = alloc_tensor({args->num_tokens, args->top_k}, dl_int32, hidden_states.device()); - - // expert_weights allocation should be done by derived class since data type could vary - - int64_t const size_of_expert_count_histogram = std::max(args->num_experts * 2, 256 * 2); - expert_count_histogram = alloc_tensor( - {size_of_expert_count_histogram}, - dl_int32, // 256 is the max number of threads per block - // and max number of experts - hidden_states.device()); - - int32_t max_num_ctas = tensorrt_llm::kernels::trtllmgen_moe::Routing::getMaxNumCtasInBatchDim( - args->num_tokens, args->top_k, args->num_experts, tile_tokens_dim); - - cta_idx_xy_to_batch_idx = alloc_tensor({max_num_ctas}, dl_int32, hidden_states.device()); - - cta_idx_xy_to_mn_limit = alloc_tensor({max_num_ctas}, dl_int32, hidden_states.device()); - - num_non_exiting_ctas = alloc_tensor({1}, dl_int32, hidden_states.device()); - - workspace.total_num_padded_tokens = static_cast(total_num_padded_tokens.data_ptr()); - workspace.total_max_padded_tokens = max_num_padded_tokens; - workspace.ProjUpTileN = tile_tokens_dim; - workspace.routing_expert_indexes = static_cast(expert_indexes.data_ptr()); - workspace.permuted_idx_size = static_cast(total_num_padded_tokens.data_ptr()); - workspace.expanded_idx_to_permuted_idx = static_cast(expanded_idx_to_permuted_idx.data_ptr()); - workspace.permuted_idx_to_token_idx = static_cast(permuted_idx_to_token_idx.data_ptr()); - // workspace.expert_weights will be set by derived class after expert_weights allocation - workspace.cta_idx_xy_to_batch_idx = static_cast(cta_idx_xy_to_batch_idx.data_ptr()); - workspace.cta_idx_xy_to_mn_limit = static_cast(cta_idx_xy_to_mn_limit.data_ptr()); - workspace.num_non_exiting_ctas = static_cast(num_non_exiting_ctas.data_ptr()); - - // Set dtype of score based on actual routing_logits dtype - if (routing_logits.has_value()) { - if (routing_logits.value().dtype() == dl_float32) { - mDtypeScore = btg::Dtype::Fp32; - } else { - mDtypeScore = btg::Dtype::Bfloat16; - } - } - } - - void check_moe_common() const { - // Hidden states [num_tokens, hidden_size] - TVM_FFI_ICHECK_EQ(hidden_states.ndim(), 2) << "hidden_states must be 2D."; - } - - // MoE computation phase workspace tensors (allocated in prepare_moe() or prepare_moe_common()) - Tensor gemm1_output; - Tensor activation_output; - Tensor gemm2_output; - Tensor workspace_fc1; - Tensor workspace_fc2; - Tensor output; - int64_t moe_tactic{-1}; - std::unique_ptr moe_runner; - - void prepare_moe_common(int64_t& moe_tactic) { - using RunnerType = tensorrt_llm::kernels::trtllmgen_moe::MoE::Runner; - // FIXME(siyuan): check llama4 routing after the fp4 FC1 kernels with bf16 scale factors were - // generated - bool usePerTokenScalingGemm1 = - per_token_scales.has_value() /* || - static_cast(this->routing_method_type) == RoutingMethodType::Llama4*/ - ; - // FIXME(siyuan): currently only nvfp4 x nvfp4 uses per-token scaling in both FC1 and FC2 - bool usePerTokenScalingGemm2 = per_token_scales.has_value() && mDtypeAct == btg::Dtype::E2m1; - // For FP8 block-scale (E4m3 activations, E4m3 weights) with DeepSeek FP8, use the - // weights-only Runner constructor to match the original kernel path and numerics. - if (this->mDtypeAct == btg::Dtype::E4m3 && this->mDtypeWeights == btg::Dtype::E4m3 && args->mUseDeepSeekFp8) { - moe_runner = std::make_unique( - this->mDtypeWeights, - args->mUseDeepSeekFp8, - (int32_t)tile_tokens_dim, - this->use_shuffled_weight, - this->weight_layout, - usePerTokenScalingGemm1, - usePerTokenScalingGemm2, - false, - false); - } else { - moe_runner = std::make_unique( - this->mDtypeAct, - this->mDtypeWeights, - args->mUseDeepSeekFp8, - (int32_t)tile_tokens_dim, - this->activation_type, - this->use_shuffled_weight, - this->weight_layout, - usePerTokenScalingGemm1, - usePerTokenScalingGemm2); - } - - if (moe_tactic == -1) { - moe_tactic = moe_runner->getDefaultValidConfigIndex( - args->top_k, args->hidden_size, args->intermediate_size, args->local_num_experts, args->num_tokens); - } - auto valid_cfgs = moe_runner->getValidConfigIndices( - args->top_k, args->hidden_size, args->intermediate_size, args->local_num_experts, args->num_tokens); - auto valid_it = std::find(valid_cfgs.begin(), valid_cfgs.end(), moe_tactic); - FLASHINFER_CHECK( - valid_it != valid_cfgs.end(), - "Invalid MoE tactic ", - moe_tactic, - " for tile_N=", - tile_tokens_dim, - ". Number of valid tactics for this tile is ", - valid_cfgs.size(), - ". This often indicates a stale or mismatched autotuner cache entry."); - this->moe_tactic = moe_tactic; - - auto workspace_sizes = moe_runner->getWorkspaceSizeInBytes(*args, moe_tactic); - workspace_fc1 = alloc_tensor({std::get<0>(workspace_sizes)}, dl_int8, hidden_states.device()); - workspace_fc2 = alloc_tensor({std::get<1>(workspace_sizes)}, dl_int8, hidden_states.device()); - workspace.bmm1_workspace = workspace_fc1.data_ptr(); - workspace.bmm2_workspace = workspace_fc2.data_ptr(); - } - - public: - virtual void check_routing() const = 0; - virtual void prepare_routing() = 0; - virtual void check_moe() const = 0; - virtual void prepare_moe(int64_t& moe_tactic) = 0; - - // Main entry point for all the executions. - // Do initializations prior to calling this as the initializations are different for bf16, fp8 and - // fp4. The executions are non-blocking by default. - virtual Array - run(int64_t moe_tactic, - bool enable_pdl = true, - bool use_routing_scales_on_input = false, - bool use_deep_seek_fp8 = false) { - check_routing(); - prepare_routing(); - - // Execute routing - tensorrt_llm::kernels::trtllmgen_moe::Routing::Runner routing_runner(tile_tokens_dim); - cudaStream_t routing_stream = get_stream(hidden_states.device()); - - // This base class only supports Mode 1 (FromLogits) - compute routing from logits - int32_t* expert_ids_param = nullptr; - - int16_t* replay_ptr = nullptr; - if (routing_replay_out.has_value()) { - replay_ptr = reinterpret_cast(routing_replay_out.value().data_ptr()); - } - - routing_runner.run( - args->routing_logits, - args->routing_bias, - args->num_tokens, - args->num_experts, - args->top_k, - args->n_group, - args->topk_group, - args->local_expert_offset, - args->local_num_experts, - args->routed_scaling_factor, - workspace.routing_expert_indexes, - static_cast(expert_count_histogram.data_ptr()), - static_cast(total_num_padded_tokens.data_ptr()), - static_cast(expanded_idx_to_permuted_idx.data_ptr()), - nullptr /*permuted_idx_to_expanded_idx.data_ptr()*/, - static_cast(permuted_idx_to_token_idx.data_ptr()), - expert_ids_param, - workspace.expert_weights, - static_cast(num_tokens_per_expert.data_ptr()), - static_cast(cta_idx_xy_to_batch_idx.data_ptr()), - static_cast(cta_idx_xy_to_mn_limit.data_ptr()), - static_cast(num_non_exiting_ctas.data_ptr()), - args->mDtypeElt, - mRoutingBiasDtype, - use_routing_scales_on_input, - use_deep_seek_fp8, - static_cast(routing_method_type), - routing_stream, - mRoutingLogitsDtype, - norm_topk_prob, - replay_ptr); - - check_moe(); - prepare_moe(moe_tactic); - - cudaStream_t moe_stream = get_stream(hidden_states.device()); - moe_runner->run(*args, workspace, hidden_states.device().device_id, moe_stream, moe_tactic, enable_pdl); - - if (args->do_finalize) { - return {output}; - } - return {gemm2_output, FusedMoeLauncher::expert_weights, expanded_idx_to_permuted_idx}; - } -}; - -void FusedMoeLauncher::init_common( - std::unique_ptr&& args, - int64_t tile_tokens_dim, - int64_t routing_method_type, - bool use_shuffled_weight, - int64_t weight_layout, - ActivationType activation_type, - bool norm_topk_prob) { - // Check devicearchitecture: Blackwell (SM 10.x) required - auto device = hidden_states.device().device_id; - int major = 0, minor = 0; - cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device); - cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device); - TVM_FFI_ICHECK(major == 10 || major == 12) - << "MoE kernel requires SM 10.x or SM 12.x architecture. Current device has SM " << major << minor; - this->device_version = std::make_tuple(major, minor); - - args->routing_logits = routing_logits.has_value() ? routing_logits.value().data_ptr() : nullptr; - args->routing_bias = routing_bias.has_value() ? routing_bias.value().data_ptr() : nullptr; - args->hidden_states = hidden_states.data_ptr(); - args->gemm1_weights = gemm1_weights.data_ptr(); - args->gemm2_weights = gemm2_weights.data_ptr(); - - this->args = std::move(args); - this->tile_tokens_dim = tile_tokens_dim; - this->routing_method_type = routing_method_type; - this->use_shuffled_weight = use_shuffled_weight; - TVM_FFI_ICHECK(0 <= weight_layout && weight_layout <= 2) << "the value of weight_layout is not recognized"; - this->weight_layout = static_cast(weight_layout); - this->activation_type = activation_type; - this->intermediate_size_factor = isGatedActivation(activation_type) ? 2 : 1; - this->norm_topk_prob = norm_topk_prob; -} - -class Bf16MoeLauncher : public FusedMoeLauncher { - public: - static constexpr std::array mSupportedTileNums = {8, 16, 32, 64, 128}; - - Bf16MoeLauncher( - Optional const& routing_logits, - Optional const& routing_bias, - TensorView const& expert_indices, - TensorView const& expert_weights, - TensorView const& hidden_states, - TensorView const& gemm1_weights, - TensorView const& gemm2_weights) - : FusedMoeLauncher( - routing_logits, - routing_bias, - hidden_states, - gemm1_weights, - Optional(), - Optional(), - gemm2_weights, - Optional(), - Optional()), - expert_indices(expert_indices), - expert_weights(expert_weights) {} - - void init( - std::unique_ptr&& args, - int64_t tile_tokens_dim, - int64_t routing_method_type, - bool use_shuffled_weight, - int64_t weight_layout, - ActivationType activation_type, - bool norm_topk_prob = true) { - // Do base class init and perform common checks - FusedMoeLauncher::init_common( - std::move(args), - tile_tokens_dim, - routing_method_type, - use_shuffled_weight, - weight_layout, - activation_type, - norm_topk_prob); - } - - void check_routing() const override { - FusedMoeLauncher::check_routing_common(); - if (expert_indices.ndim() == 2 && expert_indices.size(0) > 0) { - // Pre-computed routing: expert_indices is a packed tensor - // Format: (expert_id << 16) | (weight_bf16.view(int16)) - TVM_FFI_ICHECK_EQ(expert_indices.ndim(), 2) << "expert_indices must be 2D."; - TVM_FFI_ICHECK_EQ(expert_indices.size(0), hidden_states.size(0)) - << "expert_indices and hidden_states must have same number of tokens."; - TVM_FFI_ICHECK_EQ(expert_indices.size(1), args->top_k) << "expert_indices dim1 must match top_k."; - TVM_FFI_ICHECK_EQ(expert_indices.dtype(), dl_int32) << "expert_indices must be int32."; - } - - // TODO n_group, topk_group validation? - } - - void prepare_routing() override { - FusedMoeLauncher::prepare_routing_common(); - - args->mDtypeElt = btg::Dtype::Bfloat16; - args->mUseDeepSeekFp8 = false; - - // Set expert weights dtype based on routing bias - auto const routing_bias_dtype = routing_bias.has_value() ? routing_bias.value().dtype() : dl_bfloat16; - mRoutingBiasDtype = routing_bias_dtype == dl_bfloat16 ? btg::Dtype::Bfloat16 : btg::Dtype::Fp32; - - auto const routing_logits_dtype = routing_logits.has_value() ? routing_logits.value().dtype() : dl_bfloat16; - mRoutingLogitsDtype = routing_logits_dtype == dl_float32 ? btg::Dtype::Fp32 : btg::Dtype::Bfloat16; - - // Check ndim==2 and size>0 because empty placeholder tensors may have non-null data_ptr - bool has_precomputed_indices = expert_indices.ndim() == 2 && expert_indices.size(0) > 0; - if (has_precomputed_indices) { - // Use expert_indices directly - workspace.routing_expert_indexes = static_cast(const_cast(expert_indices.data_ptr())); - } - bool has_precomputed_weights = expert_weights.ndim() == 2 && expert_weights.size(0) > 0; - if (has_precomputed_weights) { - workspace.expert_weights = const_cast(expert_weights.data_ptr()); - } else { - auto ew_dtype = mDtypeScore == btg::Dtype::Fp32 ? dl_float32 : dl_bfloat16; - FusedMoeLauncher::expert_weights = - alloc_tensor({args->num_tokens, args->top_k}, ew_dtype, hidden_states.device()); - workspace.expert_weights = FusedMoeLauncher::expert_weights.data_ptr(); - } - } - - void check_moe() const override { - FusedMoeLauncher::check_moe_common(); - - TVM_FFI_ICHECK(weight_layout == batchedGemm::gemm::MatrixLayout::BlockMajorK) - << "BF16 Moe: weight_layout must be BlockMajorK"; - check_weights_shape("gemm1"); - check_weights_shape("gemm2"); - - TVM_FFI_ICHECK_EQ(args->intermediate_size % 128, 0) << "the second dimension of weights must be a multiple of 128."; - } - - void prepare_moe(int64_t& moe_tactic) override { - FusedMoeLauncher::prepare_moe_common(moe_tactic); - - int32_t max_num_padded_tokens = workspace.total_max_padded_tokens; - gemm1_output = alloc_tensor({max_num_padded_tokens, args->intermediate_size}, dl_bfloat16, hidden_states.device()); - activation_output = - alloc_tensor({max_num_padded_tokens, args->intermediate_size}, dl_bfloat16, hidden_states.device()); - gemm2_output = alloc_tensor({max_num_padded_tokens, args->hidden_size}, dl_bfloat16, hidden_states.device()); - - workspace.hidden_states_scale_linear = nullptr; - workspace.gemm1_output = gemm1_output.data_ptr(); - workspace.gemm1_output_scale = nullptr; - workspace.activation_output = activation_output.data_ptr(); - workspace.activation_output_scale = nullptr; - workspace.gemm2_output = gemm2_output.data_ptr(); - workspace.gemm2_output_scale = nullptr; - - if (args->output == nullptr) { - output = alloc_tensor({args->num_tokens, args->hidden_size}, dl_bfloat16, hidden_states.device()); - args->output = output.data_ptr(); - } - args->output_scale = nullptr; - } - - static Array> getValidConfigs( - int64_t top_k, - int64_t hidden_size, - int64_t intermediate_size, - int64_t num_local_experts, - int64_t num_tokens, - int64_t act_type, - bool use_shuffled_weight, - int64_t weight_layout) { - Array> valid_configs; - - std::vector supported_tile_nums(mSupportedTileNums.begin(), mSupportedTileNums.end()); - std::set selected_tile_nums = - computeSelectedTileN(supported_tile_nums, num_tokens, top_k, num_local_experts); - - for (int32_t tile_N : selected_tile_nums) { - auto moe_runner = std::make_unique( - btg::Dtype::Bfloat16, // dtype_act - btg::Dtype::Bfloat16, // dtype_weights - false, // useDeepSeekFp8 - tile_N, - static_cast(act_type), - use_shuffled_weight, - static_cast(weight_layout)); - - auto cfgs = - moe_runner->getValidConfigIndices(top_k, hidden_size, intermediate_size, num_local_experts, num_tokens); - - for (auto cfg : cfgs) { - valid_configs.push_back({tile_N, cfg}); - } - } - - return valid_configs; - } - - private: - TensorView expert_weights; - TensorView expert_indices; -}; - -class Fp8PerTensorLauncher : public FusedMoeLauncher { - public: - static constexpr std::array mSupportedTileNums = {8, 16, 32, 64, 128}; - - // Constructor that passes TensorView parameters to base constructor - Fp8PerTensorLauncher( - TensorView const& routing_logits, - Optional const& routing_bias, - TensorView const& hidden_states, - TensorView const& gemm1_weights, - TensorView const& output1_scales_scalar, - TensorView const& output1_scales_gate_scalar, - TensorView const& gemm2_weights, - TensorView const& output2_scales_scalar) - : FusedMoeLauncher( - Optional(routing_logits), - routing_bias, - hidden_states, - gemm1_weights, - Optional(output1_scales_scalar), - Optional(output1_scales_gate_scalar), - gemm2_weights, - Optional(output2_scales_scalar), - Optional()), - use_routing_scales_on_input(false) {} - - void init( - std::unique_ptr&& args, - int64_t tile_tokens_dim, - int64_t routing_method_type, - bool use_shuffled_weight, - int64_t weight_layout, - bool use_routing_scales_on_input_param, - ActivationType activation_type, - bool norm_topk_prob = true) { - this->use_routing_scales_on_input = use_routing_scales_on_input_param; - - auto dtype = hidden_states.dtype(); - if (dtype == dl_float16) { - mDtypeAct = btg::Dtype::Fp16; - } else if (dtype == dl_bfloat16) { - mDtypeAct = btg::Dtype::Bfloat16; - } else if (dtype == dl_float8_e4m3fn) { - mDtypeAct = btg::Dtype::E4m3; - } else { - TVM_FFI_LOG_AND_THROW(NotImplementedError) << "Unsupported input dtype for FP8 MoE."; - } - mDtypeWeights = btg::Dtype::E4m3; - - FusedMoeLauncher::init_common( - std::move(args), - tile_tokens_dim, - routing_method_type, - use_shuffled_weight, - weight_layout, - activation_type, - norm_topk_prob); - } - - void check_routing() const override { - FusedMoeLauncher::check_routing_common(); - } - - void prepare_routing() override { - FusedMoeLauncher::prepare_routing_common(); - - auto dtype = hidden_states.dtype(); - if (dtype == dl_float16) { - args->mDtypeElt = btg::Dtype::Fp16; - } else if (dtype == dl_bfloat16) { - args->mDtypeElt = btg::Dtype::Bfloat16; - } else if (dtype == dl_float8_e4m3fn) { - args->mDtypeElt = btg::Dtype::E4m3; - } else { - TVM_FFI_LOG_AND_THROW(NotImplementedError) << "Unsupported input dtype for MoE."; - } - - args->mDtypeOut = btg::Dtype::Bfloat16; - args->mUseDeepSeekFp8 = false; - - auto const routing_bias_dtype = routing_bias.has_value() ? routing_bias.value().dtype() : dl_bfloat16; - mRoutingBiasDtype = routing_bias_dtype == dl_bfloat16 ? btg::Dtype::Bfloat16 : btg::Dtype::Fp32; - - auto const routing_logits_dtype = routing_logits.has_value() ? routing_logits.value().dtype() : dl_bfloat16; - mRoutingLogitsDtype = routing_logits_dtype == dl_float32 ? btg::Dtype::Fp32 : btg::Dtype::Bfloat16; - - auto expert_weights_dtype = mRoutingLogitsDtype == btg::Dtype::Fp32 ? dl_float32 : dl_bfloat16; - expert_weights = alloc_tensor({args->num_tokens, args->top_k}, expert_weights_dtype, hidden_states.device()); - - workspace.expert_weights = expert_weights.data_ptr(); - if (static_cast(routing_method_type) == RoutingMethodType::Llama4) { - workspace.token_scales = expert_weights.data_ptr(); // Consumed by permuteGemm1 kernel - } - } - - void check_moe() const override { - FusedMoeLauncher::check_moe_common(); - - TVM_FFI_ICHECK(output1_scales_scalar.has_value()) << "output1_scales_scalar is required for FP8 MoE"; - TVM_FFI_ICHECK_EQ(output1_scales_scalar.value().dtype(), dl_float32) << "output1_scales_scalar must be float."; - TVM_FFI_ICHECK_EQ(output1_scales_scalar.value().ndim(), 1) << "output1_scales_scalar must be 1D."; - TVM_FFI_ICHECK_EQ(output1_scales_scalar.value().size(0), args->local_num_experts) - << "output1_scales_scalar has incorrect dim 0."; - - TVM_FFI_ICHECK(output1_scales_gate_scalar.has_value()) << "output1_scales_gate_scalar is required for FP8 MoE"; - TVM_FFI_ICHECK_EQ(output1_scales_gate_scalar.value().dtype(), dl_float32) - << "output1_scales_gate_scalar must be float."; - TVM_FFI_ICHECK_EQ(output1_scales_gate_scalar.value().ndim(), 1) << "output1_scales_gate_scalar must be 1D."; - TVM_FFI_ICHECK_EQ(output1_scales_gate_scalar.value().size(0), args->local_num_experts) - << "output1_scales_gate_scalar has incorrect dim 0."; - - TVM_FFI_ICHECK(output2_scales_scalar.has_value()) << "output2_scales_scalar is required for FP8 MoE"; - TVM_FFI_ICHECK_EQ(output2_scales_scalar.value().dtype(), dl_float32) << "output2_scales_scalar must be float."; - TVM_FFI_ICHECK_EQ(output2_scales_scalar.value().ndim(), 1) << "output2_scales_scalar must be 1D."; - TVM_FFI_ICHECK_EQ(output2_scales_scalar.value().size(0), args->local_num_experts) - << "output2_scales_scalar has incorrect dim 0."; - - TVM_FFI_ICHECK( - hidden_states.dtype() == dl_float8_e4m3fn || hidden_states.dtype() == dl_float16 || - hidden_states.dtype() == dl_bfloat16) - << "FP8 MoE: hidden_states must be float8_e4m3fn, float16, or bfloat16."; - TVM_FFI_ICHECK_EQ(gemm1_weights.dtype(), dl_float8_e4m3fn) << "FP8 MoE: gemm1_weights must be float8_e4m3fn."; - TVM_FFI_ICHECK_EQ(gemm2_weights.dtype(), dl_float8_e4m3fn) << "FP8 MoE: gemm2_weights must be float8_e4m3fn."; - } - - void prepare_moe(int64_t& moe_tactic) override { - FusedMoeLauncher::prepare_moe_common(moe_tactic); - - int32_t max_num_padded_tokens_gemm1 = workspace.total_max_padded_tokens + args->num_experts; - int32_t max_num_padded_tokens_gemm2 = workspace.total_max_padded_tokens; - - gemm1_output = - alloc_tensor({max_num_padded_tokens_gemm1, 2 * args->intermediate_size}, dl_uint8, hidden_states.device()); - gemm1_output_scale = alloc_tensor( - {2 * args->intermediate_size / 128, max_num_padded_tokens_gemm1}, dl_float32, hidden_states.device()); - - activation_output = - alloc_tensor({max_num_padded_tokens_gemm1, args->intermediate_size}, dl_uint8, hidden_states.device()); - activation_output_scale = - alloc_tensor({args->intermediate_size / 128, max_num_padded_tokens_gemm1}, dl_float32, hidden_states.device()); - - gemm2_output = alloc_tensor({max_num_padded_tokens_gemm2, args->hidden_size}, dl_bfloat16, hidden_states.device()); - - workspace.hidden_states_scale_linear = nullptr; - workspace.gemm1_output = gemm1_output.data_ptr(); - workspace.gemm1_output_scale = static_cast(gemm1_output_scale.data_ptr()); - workspace.activation_output = activation_output.data_ptr(); - workspace.activation_output_scale = static_cast(activation_output_scale.data_ptr()); - workspace.gemm2_output = gemm2_output.data_ptr(); - workspace.gemm2_output_scale = nullptr; - - if (args->output == nullptr) { - output = alloc_tensor({args->num_tokens, args->hidden_size}, dl_bfloat16, hidden_states.device()); - args->output = output.data_ptr(); - } - args->output_scale = nullptr; - - // Set scale pointers - TVM_FFI_ICHECK(output1_scales_scalar.has_value()); - TVM_FFI_ICHECK(output1_scales_gate_scalar.has_value()); - TVM_FFI_ICHECK(output2_scales_scalar.has_value()); - - args->output1_scales_scalar = static_cast(output1_scales_scalar.value().data_ptr()); - args->output1_scales_gate_scalar = static_cast(output1_scales_gate_scalar.value().data_ptr()); - args->output2_scales_scalar = static_cast(output2_scales_scalar.value().data_ptr()); - } - - private: - bool use_routing_scales_on_input; - Tensor gemm1_output_scale; - Tensor activation_output_scale; - - public: - static Array> getValidConfigs( - int64_t top_k, - int64_t hidden_size, - int64_t intermediate_size, - int64_t num_local_experts, - int64_t num_tokens, - int64_t act_type, - bool use_shuffled_weight, - int64_t weight_layout, - btg::Dtype dtype_act, - btg::Dtype dtype_weights) { - Array> valid_configs; - - std::vector supported_tile_nums(mSupportedTileNums.begin(), mSupportedTileNums.end()); - std::set selected_tile_nums = - computeSelectedTileN(supported_tile_nums, num_tokens, top_k, num_local_experts); - - for (int32_t tile_N : selected_tile_nums) { - auto moe_runner = std::make_unique( - dtype_act, - dtype_weights, - false, // useDeepSeekFp8 - tile_N, - static_cast(act_type), - use_shuffled_weight, - static_cast(weight_layout), - true, // usePerTokenScalingGemm1. always true for per-tensor fp8 due to llama4 routing - false, - false, - false); - - auto cfgs = - moe_runner->getValidConfigIndices(top_k, hidden_size, intermediate_size, num_local_experts, num_tokens); - - for (auto cfg : cfgs) { - valid_configs.push_back({tile_N, cfg}); - } - } - - return valid_configs; - } -}; - -class Fp8BlockScaleLauncher : public FusedMoeLauncher { - public: - static constexpr std::array mBaseSupportedTileNums = {8, 16, 32, 64, 128}; - - static std::vector getSupportedTileNums(Fp8QuantizationType quantization_type) { - std::vector tiles(mBaseSupportedTileNums.begin(), mBaseSupportedTileNums.end()); - if (quantization_type == Fp8QuantizationType::MxFp8) { - tiles.push_back(256); - } - return tiles; - } - - Fp8BlockScaleLauncher( - Optional const& routing_logits, - Optional const& routing_bias, - TensorView const& hidden_states, - TensorView const& hidden_states_scale, - TensorView const& gemm1_weights, - TensorView const& gemm1_weights_scale, - TensorView const& gemm2_weights, - TensorView const& gemm2_weights_scale, - TensorView const& expert_indices, - TensorView const& expert_weights, - Fp8QuantizationType quantization_type, - Optional const& gate_up_lora_delta = Optional(), - Optional const& activation_lora_input = Optional()) - : FusedMoeLauncher( - routing_logits, - routing_bias, - hidden_states, - gemm1_weights, - Optional(), - Optional(), - gemm2_weights, - Optional(), - Optional()), - hidden_states_scale(hidden_states_scale), - gemm1_weights_scale(gemm1_weights_scale), - gemm2_weights_scale(gemm2_weights_scale), - expert_indices(expert_indices), - expert_weights(expert_weights), - gate_up_lora_delta(gate_up_lora_delta), - activation_lora_input(activation_lora_input), - quantization_type(quantization_type) {} - - void init( - std::unique_ptr&& args, - int64_t tile_tokens_dim, - int64_t routing_method_type, - bool use_shuffled_weight, - int64_t weight_layout, - ActivationType activation_type, - bool norm_topk_prob = true) { - if (quantization_type == Fp8QuantizationType::MxFp8) { - mDtypeAct = btg::Dtype::MxE4m3; - mDtypeWeights = btg::Dtype::MxE4m3; - } else { - mDtypeAct = btg::Dtype::E4m3; - mDtypeWeights = btg::Dtype::E4m3; - } - - auto dtype = hidden_states.dtype(); - if (dtype == dl_float16) { - args->mDtypeElt = btg::Dtype::Fp16; - } else if (dtype == dl_bfloat16) { - args->mDtypeElt = btg::Dtype::Bfloat16; - } else if (dtype == dl_float8_e4m3fn) { - args->mDtypeElt = btg::Dtype::E4m3; - } else { - TVM_FFI_LOG_AND_THROW(NotImplementedError) << "Unsupported input dtype for MoE."; - } - - // Output is always bfloat16 for FP8 block scale - args->mDtypeOut = btg::Dtype::Bfloat16; - - FusedMoeLauncher::init_common( - std::move(args), - tile_tokens_dim, - routing_method_type, - use_shuffled_weight, - weight_layout, - activation_type, - norm_topk_prob); - } - - void check_routing() const override { - // Check ndim==2 and size>0 because empty placeholder tensors may have non-null data_ptr - if (expert_indices.ndim() == 2 && expert_indices.size(0) > 0) { - // Pre-computed routing: expert_indices is a packed tensor - // Format: (expert_id << 16) | (weight_bf16.view(int16)) - TVM_FFI_ICHECK_EQ(expert_indices.ndim(), 2) << "expert_indices must be 2D."; - TVM_FFI_ICHECK_EQ(expert_indices.size(0), hidden_states.size(0)) - << "expert_indices and hidden_states must have same number of tokens."; - TVM_FFI_ICHECK_EQ(expert_indices.size(1), args->top_k) << "expert_indices dim1 must match top_k."; - TVM_FFI_ICHECK_EQ(expert_indices.dtype(), dl_int32) << "expert_indices must be int32."; - } - - FusedMoeLauncher::check_routing_common(); - - if (static_cast(routing_method_type) != RoutingMethodType::DeepSeekV3) { - TVM_FFI_ICHECK(args->n_group <= 1) << "Current routing kernel (no groups) only supports n_group <= 1"; - TVM_FFI_ICHECK(args->topk_group <= 1) << "Current routing kernel (no groups) only supports topk_group <= 1"; - } - - if (static_cast(routing_method_type) == RoutingMethodType::DeepSeekV3) { - TVM_FFI_ICHECK(args->n_group != 0) << "n_group should not be zero for DeepSeekV3 routing"; - TVM_FFI_ICHECK(args->topk_group != 0) << "if n_group is given, topk_group must be given"; - TVM_FFI_ICHECK_EQ(args->num_experts % args->n_group, 0) << "num_experts must be divisible by n_group"; - // DeepSeekV3 routing supports top_k up to: - // - 8 when num_experts <= 384 (NumKimiK2Experts) - // - 22 when num_experts > 384 (NumNemotronExperts path) - // Keep this in sync with LAUNCH_ROUTING_DEEPSEEK in trtllm_fused_moe_routing_deepseek.cu. - constexpr int32_t kNumKimiK2Experts = 384; // same as in trtllm_fused_moe_routing_deepseek.cu - int32_t max_supported_top_k = args->num_experts <= kNumKimiK2Experts ? 8 : 22; - TVM_FFI_ICHECK(args->top_k <= max_supported_top_k && args->top_k > 0) - << "Current routing kernel (with groups) only supports top_k<=" << max_supported_top_k - << " && top_k>0 for num_experts=" << args->num_experts << "."; - TVM_FFI_ICHECK(args->topk_group <= 4 && args->topk_group > 0) - << "Current routing kernel only (with groups) supports topk_group<=4 && topk_group > 0."; - TVM_FFI_ICHECK_LE(args->topk_group, args->n_group) << "n_group must not be smaller than topk_group."; - TVM_FFI_ICHECK_LT(args->top_k, (args->topk_group * args->num_experts / args->n_group)) - << "top_k must be less than total number of experts in selected groups"; - } else if ( - static_cast(routing_method_type) == RoutingMethodType::Renormalize || - static_cast(routing_method_type) == RoutingMethodType::RenormalizeNaive || - static_cast(routing_method_type) == RoutingMethodType::SigmoidRenorm || - static_cast(routing_method_type) == RoutingMethodType::Sigmoid) { - TVM_FFI_ICHECK(args->top_k <= 32 && args->top_k > 0) - << "Current routing kernel (no groups) only supports top_k<=32 && top_k>0."; - } else if (static_cast(routing_method_type) == RoutingMethodType::Llama4) { - TVM_FFI_ICHECK_EQ(args->top_k, 1) << "Current routing kernel (no groups, Llama4) only supports top_k=1."; - } - - TVM_FFI_ICHECK_EQ(args->num_experts % 4, 0) << "Routing kernel expects that num_experts must be divisible by 4"; - TVM_FFI_ICHECK_GT(args->num_experts, args->top_k) << "num_experts must be greater than top_k"; - TVM_FFI_ICHECK_LE(args->local_num_experts + args->local_expert_offset, args->num_experts) - << "num_experts must be greater or equal to local_num_experts + local_expert_offset"; - } - - void prepare_routing() override { - FusedMoeLauncher::prepare_routing_common(); - - auto dtype = hidden_states.dtype(); - if (dtype == dl_float16) { - args->mDtypeElt = btg::Dtype::Fp16; - } else if (dtype == dl_bfloat16) { - args->mDtypeElt = btg::Dtype::Bfloat16; - } else if (dtype == dl_float8_e4m3fn) { - args->mDtypeElt = btg::Dtype::E4m3; - } else { - TVM_FFI_LOG_AND_THROW(NotImplementedError) << "Unsupported input dtype for MoE."; - } - - args->mUseDeepSeekFp8 = quantization_type == Fp8QuantizationType::DeepSeekFp8; - // Check ndim==2 and size>0 because empty placeholder tensors may have non-null data_ptr - bool has_precomputed_indices = expert_indices.ndim() == 2 && expert_indices.size(0) > 0; - if (has_precomputed_indices) { - // Use expert_indices directly - workspace.routing_expert_indexes = static_cast(const_cast(expert_indices.data_ptr())); - } else { - // Use routing_logits directly - args->routing_logits = static_cast(routing_logits.value().data_ptr()); - } - // Set expert weights dtype based on routing bias - auto const routing_bias_dtype = routing_bias.has_value() ? routing_bias.value().dtype() : dl_bfloat16; - mRoutingBiasDtype = routing_bias_dtype == dl_bfloat16 ? btg::Dtype::Bfloat16 : btg::Dtype::Fp32; - - auto const routing_logits_dtype = routing_logits.has_value() ? routing_logits.value().dtype() : dl_bfloat16; - mRoutingLogitsDtype = routing_logits_dtype == dl_float32 ? btg::Dtype::Fp32 : btg::Dtype::Bfloat16; - - // Check ndim==2 and size>0 because empty placeholder tensors may have non-null data_ptr - bool has_precomputed_weights = expert_weights.ndim() == 2 && expert_weights.size(0) > 0; - if (!has_precomputed_weights) { - auto ew_dtype = mDtypeScore == btg::Dtype::Fp32 ? dl_float32 : dl_bfloat16; - FusedMoeLauncher::expert_weights = - alloc_tensor({args->num_tokens, args->top_k}, ew_dtype, hidden_states.device()); - workspace.expert_weights = FusedMoeLauncher::expert_weights.data_ptr(); - } else { - workspace.expert_weights = const_cast(expert_weights.data_ptr()); - } - } - - void check_moe() const override { - FusedMoeLauncher::check_moe_common(); - - TVM_FFI_ICHECK_EQ(hidden_states.dtype(), dl_float8_e4m3fn) << "hidden_states must be fp8."; - if (quantization_type == Fp8QuantizationType::DeepSeekFp8) { - TVM_FFI_ICHECK_EQ(hidden_states_scale.dtype(), dl_float32) << "hidden_states_scale must be float."; - TVM_FFI_ICHECK_EQ(hidden_states_scale.ndim(), 2) << "hidden_states_scale must be 2D."; - TVM_FFI_ICHECK_EQ(hidden_states_scale.size(0), hidden_states.size(1) / 128) - << "hidden_states_scale dim0 must match hidden_states dim1 / 128."; - TVM_FFI_ICHECK_EQ(hidden_states_scale.size(1), args->num_tokens) - << "hidden_states_scale dim1 must match num_tokens."; - } else if (quantization_type == Fp8QuantizationType::MxFp8) { - TVM_FFI_CHECK( - weight_layout == batchedGemm::gemm::MatrixLayout::MajorK, "weight_layout must be MajorK for MxFp8."); - TVM_FFI_ICHECK_EQ(hidden_states_scale.dtype(), dl_uint8); - } else { - TVM_FFI_LOG_AND_THROW(NotImplementedError) << "trtllm_fp8_block_scale_moe only supports DeepSeekFp8 or MxFp8."; - } - - TVM_FFI_ICHECK_EQ(gemm1_weights.dtype(), dl_float8_e4m3fn) << "gemm1_weights must be fp8."; - TVM_FFI_ICHECK_EQ(gemm2_weights.dtype(), dl_float8_e4m3fn) << "gemm2_weights must be fp8."; - - if (quantization_type == Fp8QuantizationType::DeepSeekFp8) { - TVM_FFI_ICHECK_EQ(gemm1_weights_scale.dtype(), dl_float32) << "gemm1_weights_scale must be float."; - TVM_FFI_ICHECK_EQ(gemm1_weights_scale.ndim(), 3) << "gemm1_weights_scale must be 3D."; - TVM_FFI_ICHECK_EQ(gemm1_weights_scale.size(0), args->local_num_experts) - << "gemm1_weights_scale has incorrect shape."; - TVM_FFI_ICHECK_EQ(args->intermediate_size % 128, 0) << "intermediate_size must be a multiple of 128."; - TVM_FFI_ICHECK_EQ(gemm1_weights_scale.size(1), intermediate_size_factor * args->intermediate_size / 128) - << "gemm1_weights_scale has incorrect shape."; - TVM_FFI_ICHECK_EQ(gemm1_weights_scale.size(2), args->hidden_size / 128) - << "gemm1_weights_scale has incorrect shape."; - } else if (quantization_type == Fp8QuantizationType::MxFp8) { - TVM_FFI_ICHECK_EQ(gemm1_weights_scale.dtype(), dl_uint8) << "gemm1_weights_scale must be uint8."; - } else { - TVM_FFI_LOG_AND_THROW(NotImplementedError) << "trtllm_fp8_block_scale_moe only supports DeepSeekFp8 or MxFp8."; - } - - if (quantization_type == Fp8QuantizationType::DeepSeekFp8) { - TVM_FFI_ICHECK_EQ(gemm2_weights_scale.dtype(), dl_float32) << "gemm2_weights_scale must be float."; - TVM_FFI_ICHECK_EQ(gemm2_weights_scale.ndim(), 3) << "gemm2_weights_scale must be 3D."; - TVM_FFI_ICHECK_EQ(gemm2_weights_scale.size(0), args->local_num_experts) - << "gemm2_weights_scale has incorrect shape."; - TVM_FFI_ICHECK_EQ(gemm2_weights_scale.size(1), args->hidden_size / 128) - << "gemm2_weights_scale has incorrect shape."; - TVM_FFI_ICHECK_EQ(gemm2_weights_scale.size(2), args->intermediate_size / 128) - << "gemm2_weights_scale has incorrect shape."; - } else if (quantization_type == Fp8QuantizationType::MxFp8) { - TVM_FFI_ICHECK_EQ(gemm2_weights_scale.dtype(), dl_uint8) << "gemm2_weights_scale must be uint8."; - } else { - TVM_FFI_LOG_AND_THROW(NotImplementedError) << "trtllm_fp8_block_scale_moe only supports DeepSeekFp8 or MxFp8."; - } - - check_weights_shape("gemm1"); - check_weights_shape("gemm2"); - - if (gate_up_lora_delta.has_value()) { - TVM_FFI_ICHECK_EQ(gate_up_lora_delta.value().dtype(), dl_bfloat16) << "gate_up_lora_delta must be bf16."; - TVM_FFI_ICHECK_EQ(gate_up_lora_delta.value().ndim(), 3) - << "gate_up_lora_delta must be [num_tokens, top_k, 2 * intermediate_size]."; - TVM_FFI_ICHECK_EQ(gate_up_lora_delta.value().size(0), args->num_tokens); - TVM_FFI_ICHECK_EQ(gate_up_lora_delta.value().size(1), args->top_k); - TVM_FFI_ICHECK_EQ(gate_up_lora_delta.value().size(2), args->intermediate_size * intermediate_size_factor); - } - if (activation_lora_input.has_value()) { - TVM_FFI_ICHECK_EQ(activation_lora_input.value().dtype(), dl_bfloat16) << "activation_lora_input must be bf16."; - TVM_FFI_ICHECK_EQ(activation_lora_input.value().ndim(), 3) - << "activation_lora_input must be [num_tokens, top_k, intermediate_size]."; - TVM_FFI_ICHECK_EQ(activation_lora_input.value().size(0), args->num_tokens); - TVM_FFI_ICHECK_EQ(activation_lora_input.value().size(1), args->top_k); - TVM_FFI_ICHECK_EQ(activation_lora_input.value().size(2), args->intermediate_size); - } - - if (quantization_type == Fp8QuantizationType::DeepSeekFp8) { - TVM_FFI_ICHECK_EQ(args->intermediate_size % 128, 0) << "intermediate_size must be a multiple of 128."; - } - } - - void prepare_moe(int64_t& moe_tactic) override { - FusedMoeLauncher::prepare_moe_common(moe_tactic); - - // Calculate max_num_padded_tokens for gemm1 and gemm2 using maybeGetMinTokenCount - int32_t max_num_padded_tokens_gemm1 = tensorrt_llm::kernels::trtllmgen_moe::Routing::maybeGetMinTokenCount( - workspace.total_max_padded_tokens, args->intermediate_size, btg::dtypeGetNumBits(args->mDtypeElt)); - int32_t max_num_padded_tokens_gemm2 = tensorrt_llm::kernels::trtllmgen_moe::Routing::maybeGetMinTokenCount( - workspace.total_max_padded_tokens, args->hidden_size, btg::dtypeGetNumBits(args->mDtypeOut)); - - gemm1_output = alloc_tensor( - {max_num_padded_tokens_gemm1, intermediate_size_factor * args->intermediate_size}, - dl_uint8, - hidden_states.device()); - - if (quantization_type == Fp8QuantizationType::DeepSeekFp8) { - gemm1_output_scale = alloc_tensor( - {intermediate_size_factor * args->intermediate_size / 128, workspace.total_max_padded_tokens}, - dl_float32, - hidden_states.device()); - } else if (quantization_type == Fp8QuantizationType::MxFp8) { - // MxFP8 fuses the activation so no need for intermediate_size_factor - int64_t sf_size = - tensorrt_llm::computeSwizzledLayoutSFSize(max_num_padded_tokens_gemm1, args->intermediate_size / 32); - gemm1_output_scale = alloc_tensor({sf_size}, dl_uint8, hidden_states.device()); - } - - if (quantization_type == Fp8QuantizationType::DeepSeekFp8) { - activation_output = - alloc_tensor({max_num_padded_tokens_gemm1, args->intermediate_size}, dl_uint8, hidden_states.device()); - activation_output_scale = alloc_tensor( - {args->intermediate_size / 128, max_num_padded_tokens_gemm1}, dl_float32, hidden_states.device()); - } - - gemm2_output = alloc_tensor({max_num_padded_tokens_gemm2, args->hidden_size}, dl_bfloat16, hidden_states.device()); - - workspace.hidden_states_scale_linear = nullptr; - workspace.gemm1_output = gemm1_output.data_ptr(); - workspace.gemm1_output_scale = static_cast(gemm1_output_scale.data_ptr()); - if (quantization_type == Fp8QuantizationType::DeepSeekFp8) { - workspace.activation_output = activation_output.data_ptr(); - workspace.activation_output_scale = static_cast(activation_output_scale.data_ptr()); - } - workspace.gemm2_output = gemm2_output.data_ptr(); - workspace.gemm2_output_scale = nullptr; - - if (args->output == nullptr) { - output = alloc_tensor({args->num_tokens, args->hidden_size}, dl_bfloat16, hidden_states.device()); - args->output = output.data_ptr(); - } - args->output_scale = nullptr; - - args->hidden_states_scale = static_cast(hidden_states_scale.data_ptr()); - args->gemm1_weights_scale = static_cast(gemm1_weights_scale.data_ptr()); - args->gemm2_weights_scale = static_cast(gemm2_weights_scale.data_ptr()); - args->gate_up_lora_delta = gate_up_lora_delta.has_value() ? gate_up_lora_delta.value().data_ptr() : nullptr; - args->activation_lora_input = - activation_lora_input.has_value() ? activation_lora_input.value().data_ptr() : nullptr; - } - - private: - TensorView hidden_states_scale; - TensorView gemm1_weights_scale; - TensorView gemm2_weights_scale; - Tensor gemm1_output_scale; - Tensor activation_output_scale; - TensorView expert_indices; - TensorView expert_weights; - Optional gate_up_lora_delta; - Optional activation_lora_input; - Fp8QuantizationType quantization_type; - - public: - // Override to handle pre-computed routing - Array - run(int64_t moe_tactic, - bool enable_pdl = true, - bool use_routing_scales_on_input = false, - bool use_deep_seek_fp8 = false) override { - check_routing(); - prepare_routing(); - - cudaStream_t routing_stream = get_stream(hidden_states.device()); - tensorrt_llm::kernels::trtllmgen_moe::Routing::Runner routing_runner(tile_tokens_dim); - - // Check ndim==2 and size>0 because empty placeholder tensors may have non-null data_ptr - bool use_precomputed = expert_indices.ndim() == 2 && expert_indices.size(0) > 0; - // When using pre-computed routing, pass nullptr as routing_logits to tell the - // routing runner to use the pre-computed expert indices from workspace.routing_expert_indexes - // FP8 only supports Mode 1 (FromLogits) and Mode 2 (PackedPrecomputed), so expertIds is nullptr - int16_t* replay_ptr = nullptr; - if (routing_replay_out.has_value()) { - replay_ptr = reinterpret_cast(routing_replay_out.value().data_ptr()); - } - - routing_runner.run( - use_precomputed ? nullptr : args->routing_logits, - args->routing_bias, - args->num_tokens, - args->num_experts, - args->top_k, - args->n_group, - args->topk_group, - args->local_expert_offset, - args->local_num_experts, - args->routed_scaling_factor, - workspace.routing_expert_indexes, - static_cast(expert_count_histogram.data_ptr()), - static_cast(total_num_padded_tokens.data_ptr()), - static_cast(expanded_idx_to_permuted_idx.data_ptr()), - nullptr /*permuted_idx_to_expanded_idx.data_ptr()*/, - static_cast(permuted_idx_to_token_idx.data_ptr()), - nullptr, // expertIds - FP8 doesn't support UnpackedPrecomputed mode - workspace.expert_weights, - static_cast(num_tokens_per_expert.data_ptr()), - static_cast(cta_idx_xy_to_batch_idx.data_ptr()), - static_cast(cta_idx_xy_to_mn_limit.data_ptr()), - static_cast(num_non_exiting_ctas.data_ptr()), - args->mDtypeElt, - mRoutingBiasDtype, - use_routing_scales_on_input, - use_deep_seek_fp8, - static_cast(routing_method_type), - routing_stream, - mRoutingLogitsDtype, - norm_topk_prob, - replay_ptr); - - check_moe(); - prepare_moe(moe_tactic); - - cudaStream_t moe_stream = get_stream(hidden_states.device()); - moe_runner->run(*args, workspace, hidden_states.device().device_id, moe_stream, moe_tactic, enable_pdl); - - if (args->do_finalize) { - return {output}; - } - return {gemm2_output, FusedMoeLauncher::expert_weights, expanded_idx_to_permuted_idx}; - } - - static Array> getValidConfigs( - int64_t top_k, - int64_t hidden_size, - int64_t intermediate_size, - int64_t num_local_experts, - int64_t num_tokens, - bool use_shuffled_weight, - int64_t weight_layout, - btg::Dtype dtype_act, - btg::Dtype dtype_weights, - Fp8QuantizationType quantization_type, - int64_t act_type) { - Array> valid_configs; - auto activation_type = validateAndCastActivationType(act_type); - - auto supported_tile_nums = getSupportedTileNums(quantization_type); - std::set selected_tile_nums = - computeSelectedTileN(supported_tile_nums, num_tokens, top_k, num_local_experts); - - for (int32_t tile_N : selected_tile_nums) { - std::unique_ptr moe_runner; - // Keep getValidConfigs constructor path aligned with runtime prepare_moe_common(). - // This branch is for DeepSeek FP8 (E4m3 activations + E4m3 weights). - if (quantization_type == Fp8QuantizationType::DeepSeekFp8 && dtype_act == btg::Dtype::E4m3 && - dtype_weights == btg::Dtype::E4m3) { - TVM_FFI_ICHECK(static_cast(activation_type) == static_cast(ActivationType::Swiglu)) - << "DeepSeekFp8 only supports ActivationType::Swiglu, got " << static_cast(activation_type) << "."; - moe_runner = std::make_unique( - dtype_weights, - true /* useDeepSeekFp8 */, - tile_N, - use_shuffled_weight, - static_cast(weight_layout)); - } else { - // Under current trtllm_get_valid_moe_configs() dispatch rules, this else-path is - // reached only by FP8 block-scale MXFP8 (dtype_act=dtype_weights=MxE4m3). - moe_runner = std::make_unique( - dtype_act, // dtypeAct - dtype_weights, // dtypeWeights - quantization_type == Fp8QuantizationType::DeepSeekFp8, // useDeepSeekFp8 - tile_N, - activation_type, - use_shuffled_weight, - static_cast(weight_layout)); - } - - auto cfgs = - moe_runner->getValidConfigIndices(top_k, hidden_size, intermediate_size, num_local_experts, num_tokens); - - for (auto cfg : cfgs) { - valid_configs.push_back({tile_N, cfg}); - } - } - - return valid_configs; - } -}; - -class MxInt4BlockScaleLauncher : public FusedMoeLauncher { - public: - static constexpr std::array mSupportedTileNums = {8, 16, 32, 64, 128}; - - MxInt4BlockScaleLauncher( - TensorView const& routing_logits, - Optional const& routing_bias, - TensorView const& hidden_states, - TensorView const& gemm1_weights, - TensorView const& gemm1_weights_scale, - Optional const& gemm1_alpha, - Optional const& gemm1_beta, - Optional const& gemm1_clamp_limit, - TensorView const& gemm2_weights, - TensorView const& gemm2_weights_scale) - : FusedMoeLauncher( - Optional(routing_logits), - routing_bias, - hidden_states, - gemm1_weights, - Optional(), - Optional(), - gemm2_weights, - Optional(), - Optional()), - gemm1_weights_scale(gemm1_weights_scale), - gemm2_weights_scale(gemm2_weights_scale) {} - - void init( - std::unique_ptr&& args, - int64_t tile_tokens_dim, - int64_t routing_method_type, - bool norm_topk_prob = true) { - // currently only support mxint4 x bf16 - auto dtype = hidden_states.dtype(); - if (dtype == dl_bfloat16) { - args->mDtypeElt = btg::Dtype::Bfloat16; - } else { - TVM_FFI_LOG_AND_THROW(NotImplementedError) << "Unsupported input dtype for MoE."; - } - args->mDtypeOut = btg::Dtype::Bfloat16; - - mDtypeAct = btg::Dtype::Bfloat16; - mDtypeWeights = btg::Dtype::MxInt4; - - FusedMoeLauncher::init_common( - std::move(args), - tile_tokens_dim, - routing_method_type, - /*use_shuffled_weight=*/true, - static_cast(batchedGemm::gemm::MatrixLayout::BlockMajorK), - ActivationType::Swiglu, - norm_topk_prob); - } - - void check_routing() const override { - FusedMoeLauncher::check_routing_common(); - } - - void prepare_routing() override { - FusedMoeLauncher::prepare_routing_common(); - - args->mDtypeElt = mDtypeAct; - args->mUseDeepSeekFp8 = false; - // Set expert weights dtype based on routing bias - auto const routing_bias_dtype = routing_bias.has_value() ? routing_bias.value().dtype() : dl_bfloat16; - mRoutingBiasDtype = routing_bias_dtype == dl_bfloat16 ? btg::Dtype::Bfloat16 : btg::Dtype::Fp32; - - auto const routing_logits_dtype = routing_logits.has_value() ? routing_logits.value().dtype() : dl_bfloat16; - mRoutingLogitsDtype = routing_logits_dtype == dl_float32 ? btg::Dtype::Fp32 : btg::Dtype::Bfloat16; - - auto expert_weights_dtype = mRoutingLogitsDtype == btg::Dtype::Fp32 ? dl_float32 : dl_bfloat16; - expert_weights = alloc_tensor({args->num_tokens, args->top_k}, expert_weights_dtype, hidden_states.device()); - - workspace.expert_weights = expert_weights.data_ptr(); - } - - void check_moe() const override { - TVM_FFI_ICHECK(mDtypeAct == btg::Dtype::Bfloat16) << "Only Bfloat16 is supported by MxInt4 block scale MoE"; - - TVM_FFI_ICHECK_EQ(gemm1_weights.dtype(), dl_uint8) << "gemm1_weights must be uint8."; - TVM_FFI_ICHECK_EQ(gemm1_weights_scale.dtype(), dl_bfloat16) << "gemm1_weights_scale must be bf16."; - TVM_FFI_ICHECK_EQ(gemm2_weights.dtype(), dl_uint8) << "gemm2_weights must be uint8."; - TVM_FFI_ICHECK_EQ(gemm2_weights_scale.dtype(), dl_bfloat16) << "gemm2_weights_scale must be bf16."; - } - - void prepare_moe(int64_t& moe_tactic) override { - args->hidden_states = hidden_states.data_ptr(); - args->hidden_states_scale = nullptr; - args->gemm1_weights = gemm1_weights.data_ptr(); - args->gemm1_weights_scale = gemm1_weights_scale.data_ptr(); - args->gemm1_alpha = gemm1_alpha.has_value() ? static_cast(gemm1_alpha.value().data_ptr()) : nullptr; - args->gemm1_beta = gemm1_beta.has_value() ? static_cast(gemm1_beta.value().data_ptr()) : nullptr; - args->gemm1_clamp_limit = - gemm1_clamp_limit.has_value() ? static_cast(gemm1_clamp_limit.value().data_ptr()) : nullptr; - args->gemm2_weights = gemm2_weights.data_ptr(); - args->gemm2_weights_scale = gemm2_weights_scale.data_ptr(); - args->output1_scales_scalar = nullptr; - args->output1_scales_gate_scalar = nullptr; - args->output2_scales_scalar = nullptr; - - FusedMoeLauncher::prepare_moe_common(moe_tactic); - - max_num_padded_tokens_gemm1 = tensorrt_llm::kernels::trtllmgen_moe::Routing::maybeGetMinTokenCount( - workspace.total_max_padded_tokens, args->intermediate_size, btg::dtypeGetNumBits(mDtypeAct)); - max_num_padded_tokens_gemm2 = tensorrt_llm::kernels::trtllmgen_moe::Routing::maybeGetMinTokenCount( - workspace.total_max_padded_tokens, - args->hidden_size, - btg::dtypeGetNumBits(btg::Dtype::Bfloat16)); // Output is always BF16 - - auto const gemm1_output_hidden = args->intermediate_size; - gemm1_output = - alloc_tensor({max_num_padded_tokens_gemm1, gemm1_output_hidden}, dl_bfloat16, hidden_states.device()); - - // Allocate gemm2_output - gemm2_output = alloc_tensor({max_num_padded_tokens_gemm2, args->hidden_size}, dl_bfloat16, hidden_states.device()); - - // Setup workspace pointers - workspace.hidden_states_scale_linear = nullptr; // MxInt4 doesn't use linear scale - workspace.gemm1_output = gemm1_output.data_ptr(); - workspace.gemm1_output_scale = nullptr; - // Note: activation_output and activation_output_scale are set by the base class - // prepare_moe_common() when gated activation is used - workspace.gemm2_output = gemm2_output.data_ptr(); - workspace.gemm2_output_scale = nullptr; - } - - private: - TensorView gemm1_weights_scale; - Optional gemm1_alpha; - Optional gemm1_beta; - Optional gemm1_clamp_limit; - TensorView gemm2_weights_scale; - int32_t max_num_padded_tokens_gemm1{}; - int32_t max_num_padded_tokens_gemm2{}; - - public: - static Array> getValidConfigs( - int64_t top_k, int64_t hidden_size, int64_t intermediate_size, int64_t num_local_experts, int64_t num_tokens) { - Array> valid_configs; - - std::vector tile_sizes(mSupportedTileNums.begin(), mSupportedTileNums.end()); - std::set selected_tile_nums = computeSelectedTileN(tile_sizes, num_tokens, top_k, num_local_experts); - - for (int32_t tile_N : selected_tile_nums) { - auto moe_runner = std::make_unique( - btg::Dtype::Bfloat16, - btg::Dtype::MxInt4, - false, // useDeepSeekFp8 - tile_N, - ActivationType::Swiglu, - /*useShuffledMatrix*/ true, - batchedGemm::gemm::MatrixLayout::BlockMajorK); - - auto cfgs = - moe_runner->getValidConfigIndices(top_k, hidden_size, intermediate_size, num_local_experts, num_tokens); - - for (auto cfg : cfgs) { - valid_configs.push_back({tile_N, cfg}); - } - } - - return valid_configs; - } -}; - -class FP4BlockScaleLauncher : public FusedMoeLauncher { - public: - static constexpr std::array mBaseSupportedTileNums = {8, 16, 32, 64}; - - static std::vector getSupportedTileNums(btg::Dtype dtype_act) { - std::vector tiles(mBaseSupportedTileNums.begin(), mBaseSupportedTileNums.end()); - if (dtype_act != btg::Dtype::Bfloat16) { - tiles.push_back(128); - tiles.push_back(256); - } - return tiles; - } - - FP4BlockScaleLauncher( - RoutingInputMode routing_input_mode, - Optional const& routing_logits, - Optional const& routing_bias, - TensorView const& hidden_states, - Optional const& hidden_states_scale, - TensorView const& gemm1_weights, - TensorView const& gemm1_weights_scale, - Optional const& gemm1_bias, - Optional const& gemm1_alpha, - Optional const& gemm1_beta, - Optional const& gemm1_clamp_limit, - TensorView const& gemm2_weights, - TensorView const& gemm2_weights_scale, - Optional const& gemm2_bias, - Optional const& output1_scales_scalar, - Optional const& output1_scales_gate_scalar, - Optional const& output2_scales_scalar, - Optional const& per_token_scales, - TensorView const& topk_ids, - TensorView const& topk_weights) - : FusedMoeLauncher( - routing_logits, - routing_bias, - hidden_states, - gemm1_weights, - output1_scales_scalar, - output1_scales_gate_scalar, - gemm2_weights, - output2_scales_scalar, - per_token_scales), - routing_input_mode_(routing_input_mode), - hidden_states_scale(hidden_states_scale), - gemm1_weights_scale(gemm1_weights_scale), - gemm1_bias(gemm1_bias), - gemm1_alpha(gemm1_alpha), - gemm1_beta(gemm1_beta), - gemm1_clamp_limit(gemm1_clamp_limit), - gemm2_weights_scale(gemm2_weights_scale), - gemm2_bias(gemm2_bias), - topk_ids(topk_ids), - topk_weights(topk_weights) {} - - void init( - std::unique_ptr&& args, - int64_t tile_tokens_dim, - int64_t routing_method_type, - bool use_shuffled_weight, - int64_t weight_layout, - ActivationType activation_type, - btg::Dtype dtype_act, - btg::Dtype dtype_weights, - bool norm_topk_prob = true) { - // Set data types - args->mDtypeElt = dtype_act; - args->mDtypeOut = btg::Dtype::Bfloat16; // Output is always BF16 for FP4 - args->mUseDeepSeekFp8 = false; // FP4 doesn't use DeepSeek FP8 - - mDtypeAct = dtype_act; - mDtypeWeights = dtype_weights; - - FusedMoeLauncher::init_common( - std::move(args), - tile_tokens_dim, - routing_method_type, - use_shuffled_weight, - weight_layout, - activation_type, - norm_topk_prob); - } - - void check_routing() const override { - // First call base class common routing checks - FusedMoeLauncher::check_routing_common(); - } - - void prepare_routing() override { - num_tokens_per_expert = alloc_tensor({args->num_experts}, dl_int32, hidden_states.device()); - int32_t max_num_padded_tokens = tensorrt_llm::kernels::trtllmgen_moe::Routing::getMaxPermutedPaddedCount( - args->num_tokens, args->top_k, args->num_experts, tile_tokens_dim); - - total_num_padded_tokens = alloc_tensor({1}, dl_int32, hidden_states.device()); - expanded_idx_to_permuted_idx = alloc_tensor({args->num_tokens * args->top_k}, dl_int32, hidden_states.device()); - permuted_idx_to_token_idx = alloc_tensor({max_num_padded_tokens}, dl_int32, hidden_states.device()); - - int64_t const size_of_expert_count_histogram = std::max(args->num_experts * 2, 256 * 2); - expert_count_histogram = alloc_tensor({size_of_expert_count_histogram}, dl_int32, hidden_states.device()); - - int32_t max_num_ctas = tensorrt_llm::kernels::trtllmgen_moe::Routing::getMaxNumCtasInBatchDim( - args->num_tokens, args->top_k, args->num_experts, tile_tokens_dim); - cta_idx_xy_to_batch_idx = alloc_tensor({max_num_ctas}, dl_int32, hidden_states.device()); - cta_idx_xy_to_mn_limit = alloc_tensor({max_num_ctas}, dl_int32, hidden_states.device()); - num_non_exiting_ctas = alloc_tensor({1}, dl_int32, hidden_states.device()); - - workspace.total_num_padded_tokens = static_cast(total_num_padded_tokens.data_ptr()); - workspace.total_max_padded_tokens = max_num_padded_tokens; - workspace.ProjUpTileN = tile_tokens_dim; - workspace.routing_expert_indexes = static_cast(const_cast(topk_ids.data_ptr())); - workspace.expert_weights = const_cast(topk_weights.data_ptr()); - workspace.permuted_idx_size = static_cast(total_num_padded_tokens.data_ptr()); - workspace.expanded_idx_to_permuted_idx = static_cast(expanded_idx_to_permuted_idx.data_ptr()); - workspace.permuted_idx_to_token_idx = static_cast(permuted_idx_to_token_idx.data_ptr()); - workspace.cta_idx_xy_to_batch_idx = static_cast(cta_idx_xy_to_batch_idx.data_ptr()); - workspace.cta_idx_xy_to_mn_limit = static_cast(cta_idx_xy_to_mn_limit.data_ptr()); - workspace.num_non_exiting_ctas = static_cast(num_non_exiting_ctas.data_ptr()); - - args->mDtypeElt = mDtypeAct; - auto routing_bias_dtype = routing_bias.has_value() ? routing_bias.value().dtype() : dl_bfloat16; - mRoutingBiasDtype = routing_bias_dtype == dl_bfloat16 ? btg::Dtype::Bfloat16 : btg::Dtype::Fp32; - - auto const routing_logits_dtype = routing_logits.has_value() ? routing_logits.value().dtype() : dl_bfloat16; - mRoutingLogitsDtype = routing_logits_dtype == dl_float32 ? btg::Dtype::Fp32 : btg::Dtype::Bfloat16; - } - - void check_moe() const override { - TVM_FFI_ICHECK( - mDtypeAct == btg::Dtype::E2m1 || mDtypeAct == btg::Dtype::Bfloat16 || mDtypeAct == btg::Dtype::E4m3 || - mDtypeAct == btg::Dtype::MxE4m3) - << "Only E2m1, Bfloat16, MxE4m3 and E4m3 are supported by Fp4 block scale MoE"; - - if (mDtypeAct == btg::Dtype::E2m1) { - TVM_FFI_ICHECK(mDtypeWeights == btg::Dtype::E2m1) - << "Only E2m1 and MxE2m1 are supported by block scale MoE with E2m1 activation"; - TVM_FFI_ICHECK(hidden_states_scale.has_value()) << "hidden_states_scale is required for E2m1 activation"; - TVM_FFI_ICHECK(output1_scales_scalar.has_value()) << "output1_scales_scalar is required for E2m1 activation"; - TVM_FFI_ICHECK(output1_scales_gate_scalar.has_value()) - << "output1_scales_gate_scalar is required for E2m1 activation"; - TVM_FFI_ICHECK(output2_scales_scalar.has_value()) << "output2_scales_scalar is required for E2m1 activation"; - } else if (mDtypeAct == btg::Dtype::Bfloat16 || mDtypeAct == btg::Dtype::E4m3 || mDtypeAct == btg::Dtype::MxE4m3) { - TVM_FFI_ICHECK(mDtypeWeights == btg::Dtype::MxE2m1) - << "Only MxE2m1 weights are supported by block scale MoE with Bfloat16, E4m3 or " - "MxE4m3 activation"; - } - - if (mDtypeAct == btg::Dtype::E4m3) { - TVM_FFI_ICHECK(output1_scales_scalar.has_value()) << "output1_scales_scalar is required for E4m3 activation"; - TVM_FFI_ICHECK(output1_scales_gate_scalar.has_value()) - << "output1_scales_gate_scalar is required for E4m3 activation"; - TVM_FFI_ICHECK(output2_scales_scalar.has_value()) << "output2_scales_scalar is required for E4m3 activation"; - } - - TVM_FFI_ICHECK_EQ(gemm1_weights.dtype(), dl_uint8) << "gemm1_weights must be byte."; - TVM_FFI_ICHECK_EQ(gemm1_weights_scale.dtype(), dl_float8_e4m3fn) << "gemm1_weights_scale must be fp8."; - TVM_FFI_ICHECK_EQ(gemm2_weights.dtype(), dl_uint8) << "gemm2_weights must be byte."; - TVM_FFI_ICHECK_EQ(gemm2_weights_scale.dtype(), dl_float8_e4m3fn) << "gemm2_weights_scale must be fp8."; - } - - void prepare_moe(int64_t& moe_tactic) override { - args->hidden_states = hidden_states.data_ptr(); - args->hidden_states_scale = hidden_states_scale.has_value() ? hidden_states_scale.value().data_ptr() : nullptr; - args->gemm1_weights = gemm1_weights.data_ptr(); - args->gemm1_weights_scale = gemm1_weights_scale.data_ptr(); - args->gemm1_bias = gemm1_bias.has_value() ? static_cast(gemm1_bias.value().data_ptr()) : nullptr; - args->gemm1_alpha = gemm1_alpha.has_value() ? static_cast(gemm1_alpha.value().data_ptr()) : nullptr; - args->gemm1_beta = gemm1_beta.has_value() ? static_cast(gemm1_beta.value().data_ptr()) : nullptr; - args->gemm1_clamp_limit = - gemm1_clamp_limit.has_value() ? static_cast(gemm1_clamp_limit.value().data_ptr()) : nullptr; - args->gemm2_weights = gemm2_weights.data_ptr(); - args->gemm2_weights_scale = gemm2_weights_scale.data_ptr(); - args->gemm2_bias = gemm2_bias.has_value() ? static_cast(gemm2_bias.value().data_ptr()) : nullptr; - args->output1_scales_scalar = - output1_scales_scalar.has_value() ? static_cast(output1_scales_scalar.value().data_ptr()) : nullptr; - args->output1_scales_gate_scalar = output1_scales_gate_scalar.has_value() - ? static_cast(output1_scales_gate_scalar.value().data_ptr()) - : nullptr; - args->output2_scales_scalar = - output2_scales_scalar.has_value() ? static_cast(output2_scales_scalar.value().data_ptr()) : nullptr; - - FusedMoeLauncher::prepare_moe_common(moe_tactic); - - auto const sf_vec_size = mDtypeWeights == btg::Dtype::MxE2m1 ? 32 : 16; - - max_num_padded_tokens_gemm1 = tensorrt_llm::kernels::trtllmgen_moe::Routing::maybeGetMinTokenCount( - workspace.total_max_padded_tokens, args->intermediate_size, btg::dtypeGetNumBits(mDtypeAct)); - max_num_padded_tokens_gemm2 = tensorrt_llm::kernels::trtllmgen_moe::Routing::maybeGetMinTokenCount( - workspace.total_max_padded_tokens, - args->hidden_size, - btg::dtypeGetNumBits(btg::Dtype::Bfloat16)); // Output is always BF16 - - auto const gemm1_output_hidden = - mDtypeAct == btg::Dtype::E2m1 ? args->intermediate_size / 2 : args->intermediate_size; - if (mDtypeAct == btg::Dtype::E2m1 || mDtypeAct == btg::Dtype::MxE4m3) { - int64_t sf_size = - tensorrt_llm::computeSwizzledLayoutSFSize(max_num_padded_tokens_gemm1, args->intermediate_size / sf_vec_size); - gemm1_output_scale = alloc_tensor({sf_size}, dl_uint8, hidden_states.device()); - } - if (!per_token_scales.has_value()) { - gemm1_output = alloc_tensor( - {max_num_padded_tokens_gemm1, gemm1_output_hidden}, - mDtypeAct == btg::Dtype::Bfloat16 ? dl_bfloat16 : dl_uint8, - hidden_states.device()); - } else { // FC1 output is Bfloat16 - TVM_FFI_ICHECK(mDtypeAct == btg::Dtype::E2m1) - << "NvFP4 MoE: currently only support NvFP4 x NvFP4 when using per-token scaling."; - // When per-token scales are used, the FC1 output is always BF16 and will be quantized - gemm1_output = - alloc_tensor({max_num_padded_tokens_gemm1, args->intermediate_size}, dl_bfloat16, hidden_states.device()); - activation_output = - alloc_tensor({max_num_padded_tokens_gemm1, gemm1_output_hidden}, dl_uint8, hidden_states.device()); - per_token_scales_fc2 = alloc_tensor({max_num_padded_tokens_gemm1}, dl_float32, hidden_states.device()); - } - - // Allocate gemm2_output - gemm2_output = alloc_tensor({max_num_padded_tokens_gemm2, args->hidden_size}, dl_bfloat16, hidden_states.device()); - - // Setup workspace pointers - workspace.hidden_states_scale_linear = nullptr; // FP4 doesn't use linear scale - workspace.gemm1_output = gemm1_output.data_ptr(); - workspace.gemm1_output_scale = - gemm1_output_scale.has_value() ? static_cast(gemm1_output_scale.value().data_ptr()) : nullptr; - if (per_token_scales.has_value()) { - workspace.token_scales = per_token_scales.value().data_ptr(); - workspace.activation_output = activation_output.data_ptr(); - workspace.activation_output_scale = workspace.gemm1_output_scale; - workspace.token_scales_fc2 = per_token_scales_fc2.data_ptr(); - } - workspace.gemm2_output = gemm2_output.data_ptr(); - workspace.gemm2_output_scale = nullptr; - } - - private: - RoutingInputMode routing_input_mode_; - Optional hidden_states_scale; - TensorView gemm1_weights_scale; - Optional gemm1_bias; - Optional gemm1_alpha; - Optional gemm1_beta; - Optional gemm1_clamp_limit; - TensorView gemm2_weights_scale; - Optional gemm2_bias; - int32_t max_num_padded_tokens_gemm1{}; - int32_t max_num_padded_tokens_gemm2{}; - Optional gemm1_output_scale; - TensorView topk_ids; // [num_tokens, top_k] - pre-computed or output top-k expert indices - TensorView topk_weights; // [num_tokens, top_k] - pre-computed or output top-k routing weights - - public: - Array - run(int64_t moe_tactic, - bool enable_pdl = true, - bool use_routing_scales_on_input = false, - bool use_deep_seek_fp8 = false) override { - check_routing(); - prepare_routing(); - - // Execute routing - tensorrt_llm::kernels::trtllmgen_moe::Routing::Runner routing_runner(tile_tokens_dim); - cudaStream_t routing_stream = get_stream(hidden_states.device()); - - // Set routing kernel parameters based on mode (see RoutingInputMode enum for documentation) - int32_t* expert_ids_param = nullptr; // INPUT: pre-computed expert IDs (Mode 3 only) - void* expert_weights_param = nullptr; // INPUT or OUTPUT depending on mode - - switch (routing_input_mode_) { - case RoutingInputMode::FromLogits: - // Mode 1: Kernel computes routing, writes weights to expert_weights_param (OUTPUT) - expert_ids_param = nullptr; - expert_weights_param = topk_weights.data_ptr(); - break; - - case RoutingInputMode::PackedPrecomputed: - // Mode 2: Kernel unpacks from topk_ids, writes weights to expert_weights_param (OUTPUT) - expert_ids_param = nullptr; - expert_weights_param = topk_weights.data_ptr(); - break; - - case RoutingInputMode::UnpackedPrecomputed: - // Mode 3: Both are INPUTS, kernel uses them directly - expert_ids_param = static_cast(topk_ids.data_ptr()); - expert_weights_param = topk_weights.data_ptr(); - break; - } - - int16_t* replay_ptr = nullptr; - if (routing_replay_out.has_value()) { - replay_ptr = reinterpret_cast(routing_replay_out.value().data_ptr()); - } - - routing_runner.run( - args->routing_logits, - args->routing_bias, - args->num_tokens, - args->num_experts, - args->top_k, - args->n_group, - args->topk_group, - args->local_expert_offset, - args->local_num_experts, - args->routed_scaling_factor, - static_cast(topk_ids.data_ptr()), - static_cast(expert_count_histogram.data_ptr()), - static_cast(total_num_padded_tokens.data_ptr()), - static_cast(expanded_idx_to_permuted_idx.data_ptr()), - nullptr /*permuted_idx_to_expanded_idx*/, - static_cast(permuted_idx_to_token_idx.data_ptr()), - expert_ids_param, - expert_weights_param, - static_cast(num_tokens_per_expert.data_ptr()), - static_cast(cta_idx_xy_to_batch_idx.data_ptr()), - static_cast(cta_idx_xy_to_mn_limit.data_ptr()), - static_cast(num_non_exiting_ctas.data_ptr()), - args->mDtypeElt, - mRoutingBiasDtype, - use_routing_scales_on_input, - use_deep_seek_fp8, - static_cast(routing_method_type), - routing_stream, - mRoutingLogitsDtype, - norm_topk_prob, - replay_ptr); - - check_moe(); - prepare_moe(moe_tactic); - - cudaStream_t moe_stream = get_stream(hidden_states.device()); - moe_runner->run(*args, workspace, hidden_states.device().device_id, moe_stream, moe_tactic, enable_pdl); - - // Match original FP4 behavior for return values - if (args->do_finalize) { - return {output}; - } - return {gemm2_output, FusedMoeLauncher::expert_weights, expanded_idx_to_permuted_idx}; - } - - static Array> getValidConfigs( - int64_t top_k, - int64_t hidden_size, - int64_t intermediate_size, - int64_t num_local_experts, - int64_t num_tokens, - int64_t act_type, - btg::Dtype dtype_act, - btg::Dtype dtype_weights, - bool use_per_token_scaling) { - Array> valid_configs; - - std::vector tile_sizes = getSupportedTileNums(dtype_act); - std::set selected_tile_nums = computeSelectedTileN(tile_sizes, num_tokens, top_k, num_local_experts); - - for (int32_t tile_N : selected_tile_nums) { - auto moe_runner = std::make_unique( - dtype_act, - dtype_weights, - false, // useDeepSeekFp8 - tile_N, - static_cast(act_type), - /*useShuffledMatrix*/ true, - /*weight_layout*/ batchedGemm::gemm::MatrixLayout::MajorK, - // NOTE(siyuan): currently FP4 MoE always apply per-token scaling to both FC1 and FC2. - /*usePerTokenScalingGemm1*/ use_per_token_scaling, - /*usePerTokenScalingGemm2*/ use_per_token_scaling, - false, - false); - - auto cfgs = - moe_runner->getValidConfigIndices(top_k, hidden_size, intermediate_size, num_local_experts, num_tokens); - - for (auto cfg : cfgs) { - valid_configs.push_back({tile_N, cfg}); - } - } - - return valid_configs; - } -}; - -Array trtllm_bf16_moe( - Optional const& routing_logits, - Optional const& routing_bias, - TensorView const& expert_indices, - TensorView const& expert_weights, - TensorView const& hidden_states, - TensorView const& gemm1_weights, - TensorView const& gemm2_weights, - TensorView output, - int64_t num_experts, - int64_t top_k, - Optional n_group, - Optional topk_group, - int64_t intermediate_size, - int64_t local_expert_offset, - int64_t local_num_experts, - Optional routed_scaling_factor, - int64_t routing_method_type, - bool use_shuffled_weight, - int64_t weight_layout, - bool do_finalize, - bool enable_pdl, - Array moe_tactic, - int64_t activation_type, - bool norm_topk_prob, - Optional routing_replay_out) { - // Just some basic type validation first and leave more checks to the launcher - if (routing_logits.has_value()) { - TVM_FFI_ICHECK(routing_logits.value().dtype() == dl_float32 || routing_logits.value().dtype() == dl_bfloat16) - << "BF16 MoE: routing_logits must be bfloat16 or float."; - } - TVM_FFI_ICHECK_EQ(hidden_states.dtype(), dl_bfloat16) << "BF16 MoE: hidden_states must be bfloat16."; - TVM_FFI_ICHECK_EQ(gemm1_weights.dtype(), dl_bfloat16) << "BF16 MoE: gemm1_weights must be bfloat16."; - TVM_FFI_ICHECK_EQ(gemm2_weights.dtype(), dl_bfloat16) << "BF16 MoE: gemm2_weights must be bfloat16."; - - if (routing_replay_out.has_value()) { - validate_routing_replay_out(routing_replay_out.value(), hidden_states, top_k); - } - - auto const num_tokens = hidden_states.size(0); - auto const hidden_size = hidden_states.size(1); - auto const activation = validateAndCastActivationType(activation_type); - - // Calculate supported tile sizes - std::vector mSupportedTileN( - Bf16MoeLauncher::mSupportedTileNums.begin(), Bf16MoeLauncher::mSupportedTileNums.end()); - // Build launchers for ALL supported tiles (not just the computeSelectedTileN subset) - // so that autotuner-cached tactics always find their tile_N in the map. - // Launcher creation is cheap (no GPU allocation until run()), so this is safe. - - // Create a map of launchers for each tile size - std::unordered_map> launchers_map; - - for (int32_t curr_tile_N : mSupportedTileN) { - // Create MoE arguments for this launcher - auto args = std::make_unique(); - args->num_tokens = num_tokens; - args->num_experts = num_experts; - args->hidden_size = hidden_size; - args->hidden_size_output = args->hidden_size; - args->top_k = top_k; - args->n_group = n_group.value_or(0); - args->topk_group = topk_group.value_or(0); - args->routed_scaling_factor = routed_scaling_factor.value_or(1.0); - args->local_expert_offset = local_expert_offset; - args->local_num_experts = local_num_experts; - args->intermediate_size = intermediate_size; - args->do_finalize = do_finalize; - args->output = output.data_ptr(); - args->output_scale = nullptr; - - // Create and initialize launcher for this tile size - auto launcher = std::make_unique( - routing_logits, routing_bias, expert_indices, expert_weights, hidden_states, gemm1_weights, gemm2_weights); - launcher->init( - std::move(args), - curr_tile_N, - routing_method_type, - use_shuffled_weight, - weight_layout, - activation, - norm_topk_prob); - launcher->set_routing_replay_out(routing_replay_out); - - launchers_map[curr_tile_N] = std::move(launcher); - } - - auto const [tile_N, config] = - resolveMoeTileAndConfig(moe_tactic, mSupportedTileN, num_tokens, top_k, local_num_experts); - - // Get the launcher for the selected tile_N - auto launcher_it = launchers_map.find(static_cast(tile_N)); - FLASHINFER_CHECK(launcher_it != launchers_map.end(), "Internal error: missing BF16 MoE launcher for tile_N=", tile_N); - auto& selected_launcher = launcher_it->second; - - // Run the launcher - it will create its own runner internally - return selected_launcher->run(config, enable_pdl); -} - -Array trtllm_fp8_per_tensor_scale_moe( - TensorView routing_logits, - Optional routing_bias, - TensorView hidden_states, - TensorView gemm1_weights, - TensorView output1_scales_scalar, - TensorView output1_scales_gate_scalar, - TensorView gemm2_weights, - TensorView output2_scales_scalar, - TensorView output, - int64_t num_experts, - int64_t top_k, - Optional n_group, - Optional topk_group, - int64_t intermediate_size, - int64_t local_expert_offset, - int64_t local_num_experts, - Optional routed_scaling_factor, - bool use_routing_scales_on_input, - int64_t routing_method_type, - bool do_finalize, - bool enable_pdl, - Array config_index, - int64_t activation_type, - bool norm_topk_prob, - Optional routing_replay_out) { - // Basic type validation - auto dtype = hidden_states.dtype(); - auto activation = validateAndCastActivationType(activation_type); - - TVM_FFI_ICHECK(dtype == dl_float8_e4m3fn || dtype == dl_float16 || dtype == dl_bfloat16) - << "FP8 MoE: hidden_states must be float8_e4m3fn, float16, or bfloat16."; - TVM_FFI_ICHECK_EQ(gemm1_weights.dtype(), dl_float8_e4m3fn) << "FP8 MoE: gemm1_weights must be float8_e4m3fn."; - TVM_FFI_ICHECK_EQ(gemm2_weights.dtype(), dl_float8_e4m3fn) << "FP8 MoE: gemm2_weights must be float8_e4m3fn."; - TVM_FFI_ICHECK_EQ(output1_scales_scalar.dtype(), dl_float32) << "FP8 MoE: output1_scales_scalar must be float32."; - TVM_FFI_ICHECK_EQ(output1_scales_gate_scalar.dtype(), dl_float32) - << "FP8 MoE: output1_scales_gate_scalar must be float32."; - TVM_FFI_ICHECK_EQ(output2_scales_scalar.dtype(), dl_float32) << "FP8 MoE: output2_scales_scalar must be float32."; - - if (routing_replay_out.has_value()) { - validate_routing_replay_out(routing_replay_out.value(), hidden_states, top_k); - } - - auto const num_tokens = hidden_states.size(0); - auto const hidden_size = hidden_states.size(1); - - // Use default values that match the original function behavior - bool use_shuffled_weight = true; // Original uses /*useShuffledMatrix*/ true - int64_t weight_layout = 0; // Default to MajorK - - // Calculate supported tile sizes - std::vector mSupportedTileN( - Fp8PerTensorLauncher::mSupportedTileNums.begin(), Fp8PerTensorLauncher::mSupportedTileNums.end()); - // Build launchers for ALL supported tiles so autotuner-cached tactics always find their tile_N. - - // Create a map of launchers for each tile size - std::unordered_map> launchers_map; - - for (int32_t curr_tile_N : mSupportedTileN) { - // Create MoE arguments for this launcher - auto args = std::make_unique(); - args->num_tokens = num_tokens; - args->num_experts = num_experts; - args->hidden_size = hidden_size; - args->hidden_size_output = args->hidden_size; - args->top_k = top_k; - args->n_group = n_group.value_or(0); - args->topk_group = topk_group.value_or(0); - args->local_expert_offset = local_expert_offset; - args->local_num_experts = local_num_experts; - args->intermediate_size = intermediate_size; - args->routed_scaling_factor = routed_scaling_factor.value_or(1.0); - args->do_finalize = do_finalize; - args->output = output.data_ptr(); - args->output_scale = nullptr; - - // Create and initialize launcher for this tile size - auto launcher = std::make_unique( - routing_logits, - routing_bias, - hidden_states, - gemm1_weights, - output1_scales_scalar, - output1_scales_gate_scalar, - gemm2_weights, - output2_scales_scalar); - launcher->init( - std::move(args), - curr_tile_N, - routing_method_type, - use_shuffled_weight, - weight_layout, - use_routing_scales_on_input, - activation, - norm_topk_prob); - launcher->set_routing_replay_out(routing_replay_out); - - launchers_map[curr_tile_N] = std::move(launcher); - } - - auto const [tile_N, config] = - resolveMoeTileAndConfig(config_index, mSupportedTileN, num_tokens, top_k, local_num_experts); - - // Get the launcher for the selected tile_N - auto launcher_it = launchers_map.find(static_cast(tile_N)); - FLASHINFER_CHECK( - launcher_it != launchers_map.end(), "Internal error: missing FP8 per-tensor MoE launcher for tile_N=", tile_N); - auto& selected_launcher = launcher_it->second; - - // Run the launcher - it will create its own runner internally - return selected_launcher->run(config, enable_pdl, use_routing_scales_on_input); -} - -Array trtllm_fp8_block_scale_moe_impl( - Optional routing_logits, - TensorView expert_indices, - TensorView expert_weights, - Optional routing_bias, - TensorView hidden_states, - TensorView hidden_states_scale, - TensorView gemm1_weights, - TensorView gemm1_weights_scale, - TensorView gemm2_weights, - TensorView gemm2_weights_scale, - TensorView output, - int64_t num_experts, - int64_t top_k, - Optional n_group, - Optional topk_group, - int64_t intermediate_size, - int64_t local_expert_offset, - int64_t local_num_experts, - Optional routed_scaling_factor, - int64_t routing_method_type, - bool use_shuffled_weight, - int64_t weight_layout, - bool do_finalize, - bool enable_pdl, - Array config_index, - Fp8QuantizationType quantization_type, - int64_t act_type, - bool norm_topk_prob, - Optional routing_replay_out, - Optional gate_up_lora_delta, - Optional activation_lora_input, - int64_t lora_ready_event = 0, - int64_t gemm2_done_event = 0) { - auto activation_type = validateAndCastActivationType(act_type); - // DeepSeekFp8 currently uses a TRTLLM runner that hardwires Swiglu activation semantics. - // Fail for any other activation to avoid silently running incorrect activation behavior. - if (quantization_type == Fp8QuantizationType::DeepSeekFp8 && activation_type != ActivationType::Swiglu) { - TVM_FFI_LOG_AND_THROW(NotImplementedError) - << "DeepSeekFp8 only supports ActivationType::Swiglu in this runner path. " - << "Received activation_type=" << static_cast(activation_type); - } - - // Basic type validation - auto dtype = hidden_states.dtype(); - - // Either routing_logits or expert_indices must be provided - // expert_indices is a packed tensor: (expert_id << 16) | (weight_bf16.view(int16)) - bool use_routing_logits = routing_logits.has_value(); - // Check ndim==2 and size>0 because empty placeholder tensors may have non-null data_ptr - bool use_precomputed_routing = expert_indices.ndim() == 2 && expert_indices.size(0) > 0; - - TVM_FFI_ICHECK(use_routing_logits || use_precomputed_routing) - << "Either routing_logits or expert_indices must be provided."; - - (void)use_routing_logits; - TVM_FFI_ICHECK(dtype == dl_float16 || dtype == dl_bfloat16 || dtype == dl_float8_e4m3fn) - << "FP8 block scale MoE: hidden_states must be fp16, bf16, or fp8."; - if (quantization_type == Fp8QuantizationType::DeepSeekFp8) { - TVM_FFI_ICHECK_EQ(hidden_states_scale.dtype(), dl_float32) - << "FP8 block scale MoE: hidden_states_scale must be float32."; - } else if (quantization_type == Fp8QuantizationType::MxFp8) { - TVM_FFI_ICHECK_EQ(hidden_states_scale.dtype(), dl_uint8) - << "FP8 block scale MoE: hidden_states_scale must be uint8."; - } else { - TVM_FFI_LOG_AND_THROW(NotImplementedError) << "trtllm_fp8_block_scale_moe only supports DeepSeekFp8 or MxFp8."; - } - TVM_FFI_ICHECK_EQ(gemm1_weights.dtype(), dl_float8_e4m3fn) << "FP8 block scale MoE: gemm1_weights must be fp8."; - TVM_FFI_ICHECK_EQ(gemm2_weights.dtype(), dl_float8_e4m3fn) << "FP8 block scale MoE: gemm2_weights must be fp8."; - if (quantization_type == Fp8QuantizationType::DeepSeekFp8) { - TVM_FFI_ICHECK_EQ(gemm1_weights_scale.dtype(), dl_float32) - << "FP8 block scale MoE: gemm1_weights_scale must be float32."; - TVM_FFI_ICHECK_EQ(gemm2_weights_scale.dtype(), dl_float32) - << "FP8 block scale MoE: gemm2_weights_scale must be float32."; - } else if (quantization_type == Fp8QuantizationType::MxFp8) { - TVM_FFI_ICHECK_EQ(gemm1_weights_scale.dtype(), dl_uint8) - << "FP8 block scale MoE: gemm1_weights_scale must be uint8."; - TVM_FFI_ICHECK_EQ(gemm2_weights_scale.dtype(), dl_uint8) - << "FP8 block scale MoE: gemm2_weights_scale must be uint8."; - } else { - TVM_FFI_LOG_AND_THROW(NotImplementedError) << "trtllm_fp8_block_scale_moe only supports DeepSeekFp8 or MxFp8."; - } - - if (quantization_type == Fp8QuantizationType::MxFp8) { - TVM_FFI_ICHECK(use_shuffled_weight) << "use_shuffled_weight must be true for MxFp8."; - TVM_FFI_ICHECK(weight_layout == 0) << "weight_layout must be 0 for MxFp8."; - } - - if (routing_replay_out.has_value()) { - validate_routing_replay_out(routing_replay_out.value(), hidden_states, top_k); - } - - auto const num_tokens = hidden_states.size(0); - auto const hidden_size = hidden_states.size(1); - - auto supported_tile_nums = Fp8BlockScaleLauncher::getSupportedTileNums(quantization_type); - // Build launchers for ALL supported tiles so autotuner-cached tactics always find their tile_N. - - // Create a map of launchers for each tile size - std::unordered_map> launchers_map; - - for (int32_t curr_tile_N : supported_tile_nums) { - // Create MoE arguments for this launcher - auto args = std::make_unique(); - args->num_tokens = num_tokens; - args->num_experts = num_experts; - args->hidden_size = hidden_size; - args->hidden_size_output = args->hidden_size; - args->top_k = top_k; - args->n_group = n_group.value_or(0); - args->topk_group = topk_group.value_or(0); - args->local_expert_offset = local_expert_offset; - args->local_num_experts = local_num_experts; - args->intermediate_size = intermediate_size; - args->routed_scaling_factor = routed_scaling_factor.value_or(1.0); - args->do_finalize = do_finalize; - args->output = output.data_ptr(); - args->output_scale = nullptr; - // GEMM1-LoRA overlap: cudaEvent_t handle (recorded on the LoRA side stream) the runner - // waits on right before activation; 0 = no wait (serial path). - args->lora_ready_event = reinterpret_cast(lora_ready_event); - // Down-LoRA/finalize overlap: cudaEvent_t handle the runner records right after GEMM2 - // (before finalize) so the LoRA side stream can overlap the down-proj LoRA with - // finalize; 0 = no record (serial path). - args->gemm2_done_event = reinterpret_cast(gemm2_done_event); - - // Create and initialize launcher for this tile size - auto launcher = std::make_unique( - routing_logits, - routing_bias, - hidden_states, - hidden_states_scale, - gemm1_weights, - gemm1_weights_scale, - gemm2_weights, - gemm2_weights_scale, - expert_indices, - expert_weights, - quantization_type, - gate_up_lora_delta, - activation_lora_input); - launcher->init( - std::move(args), - curr_tile_N, - routing_method_type, - use_shuffled_weight, - weight_layout, - activation_type, - norm_topk_prob); - launcher->set_routing_replay_out(routing_replay_out); - - launchers_map[curr_tile_N] = std::move(launcher); - } - - auto const [tile_N, config] = - resolveMoeTileAndConfig(config_index, supported_tile_nums, num_tokens, top_k, local_num_experts); - - // Get the launcher for the selected tile_N - auto launcher_it = launchers_map.find(static_cast(tile_N)); - FLASHINFER_CHECK( - launcher_it != launchers_map.end(), "Internal error: missing FP8 block-scale MoE launcher for tile_N=", tile_N); - auto& selected_launcher = launcher_it->second; - - // Run the launcher with DeepSeek FP8 enabled - it will create its own runner internally - return selected_launcher->run( - config, - enable_pdl, - false /* use_routing_scales_on_input */, - quantization_type == Fp8QuantizationType::DeepSeekFp8 /* use_deep_seek_fp8 */); -} - -Array trtllm_fp8_block_scale_moe( - Optional routing_logits, - TensorView expert_indices, - TensorView expert_weights, - Optional routing_bias, - TensorView hidden_states, - TensorView hidden_states_scale, - TensorView gemm1_weights, - TensorView gemm1_weights_scale, - TensorView gemm2_weights, - TensorView gemm2_weights_scale, - TensorView output, - int64_t num_experts, - int64_t top_k, - Optional n_group, - Optional topk_group, - int64_t intermediate_size, - int64_t local_expert_offset, - int64_t local_num_experts, - Optional routed_scaling_factor, - int64_t routing_method_type, - bool use_shuffled_weight, - int64_t weight_layout, - bool do_finalize, - bool enable_pdl, - Array config_index, - Fp8QuantizationType quantization_type, - int64_t act_type, - bool norm_topk_prob, - Optional routing_replay_out) { - return trtllm_fp8_block_scale_moe_impl( - routing_logits, - expert_indices, - expert_weights, - routing_bias, - hidden_states, - hidden_states_scale, - gemm1_weights, - gemm1_weights_scale, - gemm2_weights, - gemm2_weights_scale, - output, - num_experts, - top_k, - n_group, - topk_group, - intermediate_size, - local_expert_offset, - local_num_experts, - routed_scaling_factor, - routing_method_type, - use_shuffled_weight, - weight_layout, - do_finalize, - enable_pdl, - config_index, - quantization_type, - act_type, - norm_topk_prob, - routing_replay_out, - Optional(), - Optional()); -} - -Array sgl_trtllm_fp8_block_scale_moe_lora( - Optional routing_logits, - TensorView expert_indices, - TensorView expert_weights, - Optional routing_bias, - TensorView hidden_states, - TensorView hidden_states_scale, - TensorView gemm1_weights, - TensorView gemm1_weights_scale, - TensorView gemm2_weights, - TensorView gemm2_weights_scale, - TensorView output, - int64_t num_experts, - int64_t top_k, - Optional n_group, - Optional topk_group, - int64_t intermediate_size, - int64_t local_expert_offset, - int64_t local_num_experts, - Optional routed_scaling_factor, - int64_t routing_method_type, - bool use_shuffled_weight, - int64_t weight_layout, - bool do_finalize, - bool enable_pdl, - Array config_index, - Fp8QuantizationType quantization_type, - int64_t act_type, - bool norm_topk_prob, - Optional routing_replay_out, - TensorView gate_up_lora_delta, - TensorView activation_lora_input, - int64_t lora_ready_event, - int64_t gemm2_done_event) { - if (quantization_type != Fp8QuantizationType::DeepSeekFp8) { - TVM_FFI_LOG_AND_THROW(NotImplementedError) - << "sgl_trtllm_fp8_block_scale_moe_lora currently supports DeepSeekFp8 only."; - } - return trtllm_fp8_block_scale_moe_impl( - routing_logits, - expert_indices, - expert_weights, - routing_bias, - hidden_states, - hidden_states_scale, - gemm1_weights, - gemm1_weights_scale, - gemm2_weights, - gemm2_weights_scale, - output, - num_experts, - top_k, - n_group, - topk_group, - intermediate_size, - local_expert_offset, - local_num_experts, - routed_scaling_factor, - routing_method_type, - use_shuffled_weight, - weight_layout, - do_finalize, - enable_pdl, - config_index, - quantization_type, - act_type, - norm_topk_prob, - routing_replay_out, - Optional(gate_up_lora_delta), - Optional(activation_lora_input), - lora_ready_event, - gemm2_done_event); -} - -__global__ void sgl_trtllm_fp8_block_scale_moe_lora_finalize_kernel( - cutlass::bfloat16_t const* __restrict__ gemm2_output, - cutlass::bfloat16_t const* __restrict__ expert_weights, - int32_t const* __restrict__ expanded_idx_to_permuted_idx, - cutlass::bfloat16_t const* __restrict__ down_lora_delta, - cutlass::bfloat16_t* __restrict__ output, - int64_t num_tokens, - int64_t top_k, - int64_t hidden_size, - int64_t hidden_size_padded, - float routed_scaling_factor) { - for (int64_t token_idx = blockIdx.y; token_idx < num_tokens; token_idx += gridDim.y) { - for (int64_t hidden_idx = threadIdx.x + blockDim.x * blockIdx.x; hidden_idx < hidden_size; - hidden_idx += blockDim.x * gridDim.x) { - float acc = 0.0f; - float lora_acc = 0.0f; - for (int64_t k = 0; k < top_k; ++k) { - int64_t const expanded_idx = token_idx * top_k + k; - int32_t const permuted_idx = expanded_idx_to_permuted_idx[expanded_idx]; - if (permuted_idx != -1) { - float const expert_prob = static_cast(expert_weights[token_idx * top_k + k]); - acc += expert_prob * static_cast(gemm2_output[permuted_idx * hidden_size_padded + hidden_idx]); - } - lora_acc += static_cast(down_lora_delta[expanded_idx * hidden_size + hidden_idx]); - } - output[token_idx * hidden_size + hidden_idx] = - static_cast(acc + routed_scaling_factor * lora_acc); - } - } -} - -void sgl_trtllm_fp8_block_scale_moe_lora_finalize( - TensorView gemm2_output, - TensorView expert_weights, - TensorView expanded_idx_to_permuted_idx, - TensorView down_lora_delta, - TensorView output, - Optional routed_scaling_factor) { - TVM_FFI_ICHECK_EQ(gemm2_output.dtype(), dl_bfloat16) << "gemm2_output must be bfloat16."; - TVM_FFI_ICHECK_EQ(expert_weights.dtype(), dl_bfloat16) << "expert_weights must be bfloat16."; - TVM_FFI_ICHECK((expanded_idx_to_permuted_idx.dtype() == DLDataType{kDLInt, 32, 1})) - << "expanded_idx_to_permuted_idx must be int32."; - TVM_FFI_ICHECK_EQ(down_lora_delta.dtype(), dl_bfloat16) << "down_lora_delta must be bfloat16."; - TVM_FFI_ICHECK_EQ(output.dtype(), dl_bfloat16) << "output must be bfloat16."; - TVM_FFI_ICHECK_EQ(gemm2_output.ndim(), 2) << "gemm2_output must be 2D."; - TVM_FFI_ICHECK_EQ(expert_weights.ndim(), 2) << "expert_weights must be 2D."; - TVM_FFI_ICHECK_EQ(expanded_idx_to_permuted_idx.ndim(), 1) << "expanded_idx_to_permuted_idx must be 1D."; - TVM_FFI_ICHECK_EQ(down_lora_delta.ndim(), 3) << "down_lora_delta must be 3D."; - TVM_FFI_ICHECK_EQ(output.ndim(), 2) << "output must be 2D."; - TVM_FFI_ICHECK(gemm2_output.IsContiguous()) << "gemm2_output must be contiguous."; - TVM_FFI_ICHECK(expert_weights.IsContiguous()) << "expert_weights must be contiguous."; - TVM_FFI_ICHECK(expanded_idx_to_permuted_idx.IsContiguous()) << "expanded_idx_to_permuted_idx must be contiguous."; - TVM_FFI_ICHECK(down_lora_delta.IsContiguous()) << "down_lora_delta must be contiguous."; - TVM_FFI_ICHECK(output.IsContiguous()) << "output must be contiguous."; - - int64_t const num_tokens = output.size(0); - int64_t const hidden_size = output.size(1); - int64_t const top_k = down_lora_delta.size(1); - TVM_FFI_ICHECK_EQ(expert_weights.size(0), num_tokens) << "expert_weights dim0 must equal num_tokens."; - TVM_FFI_ICHECK_EQ(expert_weights.size(1), top_k) << "expert_weights dim1 must equal top_k."; - TVM_FFI_ICHECK_EQ(down_lora_delta.size(0), num_tokens) << "down_lora_delta dim0 must equal num_tokens."; - TVM_FFI_ICHECK_EQ(down_lora_delta.size(2), hidden_size) << "down_lora_delta dim2 must equal hidden_size."; - TVM_FFI_ICHECK_EQ(expanded_idx_to_permuted_idx.size(0), num_tokens * top_k) - << "expanded_idx_to_permuted_idx size must equal num_tokens * top_k."; - TVM_FFI_ICHECK(gemm2_output.size(1) >= hidden_size) - << "gemm2_output hidden dimension is smaller than output hidden dimension."; - - int const num_threads = 128; - int const num_blocks_x = (hidden_size + num_threads - 1) / num_threads; - int const num_blocks_y = std::min(8192, num_tokens); - dim3 grid(num_blocks_x, num_blocks_y); - cudaStream_t stream = get_stream(output.device()); - sgl_trtllm_fp8_block_scale_moe_lora_finalize_kernel<<>>( - static_cast(gemm2_output.data_ptr()), - static_cast(expert_weights.data_ptr()), - static_cast(expanded_idx_to_permuted_idx.data_ptr()), - static_cast(down_lora_delta.data_ptr()), - static_cast(output.data_ptr()), - num_tokens, - top_k, - hidden_size, - gemm2_output.size(1), - static_cast(routed_scaling_factor.value_or(1.0))); - auto err = cudaGetLastError(); - FLASHINFER_CHECK(err == cudaSuccess, cudaGetErrorString(err)); -} - -Array trtllm_fp4_block_scale_moe( - int64_t routing_input_mode, - Optional routing_logits, - TensorView topk_ids, - TensorView topk_weights, - Optional routing_bias, - TensorView hidden_states, - Optional hidden_states_scale, - TensorView gemm1_weights, - TensorView gemm1_weights_scale, - Optional gemm1_bias, - Optional gemm1_alpha, - Optional gemm1_beta, - Optional gemm1_clamp_limit, - TensorView gemm2_weights, - TensorView gemm2_weights_scale, - Optional gemm2_bias, - Optional output1_scales_scalar, - Optional output1_scales_gate_scalar, - Optional output2_scales_scalar, - Optional per_token_scales, - int64_t num_experts, - int64_t top_k, - Optional n_group, - Optional topk_group, - int64_t intermediate_size, - int64_t local_expert_offset, - int64_t local_num_experts, - Optional routed_scaling_factor, - int64_t routing_method_type, - bool do_finalize, - bool enable_pdl, - int64_t act_type, - TensorView output, - Array config_index, - bool norm_topk_prob, - Optional routing_replay_out) { - // Determine data types based on input format - int const num_tokens = hidden_states.size(0); - int hidden_size = hidden_states.size(1); - if (hidden_states.dtype() == dl_uint8) hidden_size *= 2; - - int64_t hidden_states_scale_vec_size = -1; - if (hidden_states_scale.has_value()) { - hidden_states_scale_vec_size = - (static_cast(num_tokens) * hidden_size) / hidden_states_scale.value().numel(); - } - int64_t intermediate_size_factor = isGatedActivation(static_cast(act_type)) ? 2 : 1; - int64_t logical_scale_count = - static_cast(local_num_experts) * intermediate_size * intermediate_size_factor * hidden_size; - int64_t weight_scale_vec_size_raw = logical_scale_count / gemm1_weights_scale.numel(); - - // Snap to nearest valid sf_vec_size (16 or 32). - // The raw value may be slightly smaller than the true vec_size because - // block_scale_interleave pads scale columns to a multiple of 4, inflating numel(). - int64_t weight_scale_vec_size = weight_scale_vec_size_raw > 16 ? 32 : 16; - - // Round-trip validation: the unpadded scale count must not exceed actual numel - // (padding only adds elements, never removes them). - int64_t expected_unpadded = logical_scale_count / weight_scale_vec_size; - TVM_FFI_ICHECK(gemm1_weights_scale.numel() >= expected_unpadded) - << "weight scale tensor too small: numel=" << gemm1_weights_scale.numel() << " but expected at least " - << expected_unpadded << " for sf_vec_size=" << weight_scale_vec_size; - - auto mDtypeWeights = weight_scale_vec_size == 16 ? btg::Dtype::E2m1 : btg::Dtype::MxE2m1; - - if (routing_bias.has_value()) { - TVM_FFI_ICHECK(routing_bias.value().dtype() == dl_bfloat16 || routing_bias.value().dtype() == dl_float32) - << "routing_bias must be bfloat16 or float."; - - TVM_FFI_ICHECK_EQ(routing_bias.value().ndim(), 1) << "routing_bias must be 1D."; - TVM_FFI_ICHECK_EQ(routing_bias.value().size(0), num_experts) << "routing_bias has incorrect shape."; - } - - if (routing_replay_out.has_value()) { - validate_routing_replay_out(routing_replay_out.value(), hidden_states, top_k); - } - - // Determine activation type - TVM_FFI_ICHECK(gemm1_weights.dtype() == dl_uint8 && gemm2_weights.dtype() == dl_uint8) - << "weights must be fp4 packed in uint8."; - TVM_FFI_ICHECK( - hidden_states.dtype() == dl_uint8 || hidden_states.dtype() == dl_bfloat16 || - hidden_states.dtype() == dl_float8_e4m3fn) - << "hidden_states must be bf16, fp8 or uint8 (packed fp4)."; - - auto mDtypeAct = btg::Dtype::Bfloat16; - if (hidden_states.dtype() == dl_uint8) { - TVM_FFI_ICHECK(hidden_states_scale.has_value() && hidden_states_scale.value().dtype() == dl_float8_e4m3fn) - << "hidden_states_scale must be provided for fp4 activation."; - if (hidden_states_scale_vec_size == 16) { - mDtypeAct = btg::Dtype::E2m1; - } else if (hidden_states_scale_vec_size == 32) { - mDtypeAct = btg::Dtype::MxE2m1; - } else { - TVM_FFI_LOG_AND_THROW(NotImplementedError) << "Unsupported hidden state scale shape."; - } - } else if (hidden_states.dtype() == dl_float8_e4m3fn) { - if (hidden_states_scale.has_value()) { - if (hidden_states_scale_vec_size == 32) { - mDtypeAct = btg::Dtype::MxE4m3; - } else { - TVM_FFI_LOG_AND_THROW(NotImplementedError) << "Unsupported hidden state scale shape."; - } - } else { - mDtypeAct = btg::Dtype::E4m3; - } - } - - // Determine supported tile sizes - std::vector mSupportedTileN = FP4BlockScaleLauncher::getSupportedTileNums(mDtypeAct); - // Build launchers for ALL supported tiles so autotuner-cached tactics always find their tile_N. - - // Create a map of launchers for each tile size - std::unordered_map> launchers_map; - - for (int32_t curr_tile_N : mSupportedTileN) { - // Create MoE arguments for this launcher - auto args = std::make_unique(); - args->num_tokens = num_tokens; - args->num_experts = num_experts; - // For E2m1, hidden_size is already multiplied by 2 above, so use it directly - args->hidden_size = hidden_size; - args->hidden_size_output = output.size(1); - args->top_k = top_k; - args->n_group = n_group.value_or(0); - args->topk_group = topk_group.value_or(0); - args->local_expert_offset = local_expert_offset; - args->local_num_experts = local_num_experts; - args->intermediate_size = intermediate_size; - args->routed_scaling_factor = routed_scaling_factor.value_or(1.0); - args->do_finalize = do_finalize; - args->output = output.data_ptr(); - args->output_scale = nullptr; - - // Create and initialize launcher for this tile size - auto launcher = std::make_unique( - static_cast(routing_input_mode), - routing_logits, - routing_bias, - hidden_states, - hidden_states_scale, - gemm1_weights, - gemm1_weights_scale, - gemm1_bias, - gemm1_alpha, - gemm1_beta, - gemm1_clamp_limit, - gemm2_weights, - gemm2_weights_scale, - gemm2_bias, - output1_scales_scalar, - output1_scales_gate_scalar, - output2_scales_scalar, - per_token_scales, - topk_ids, - topk_weights); - launcher->init( - std::move(args), - curr_tile_N, - routing_method_type, - /*use_shuffled_weight=*/true, - /*weight_layout=*/0, - static_cast(act_type), - mDtypeAct, - mDtypeWeights, - norm_topk_prob); - launcher->set_routing_replay_out(routing_replay_out); - - launchers_map[curr_tile_N] = std::move(launcher); - } - - auto const [tile_N, config] = - resolveMoeTileAndConfig(config_index, mSupportedTileN, num_tokens, top_k, local_num_experts); - - // Get the launcher for the selected tile_N - auto launcher_it = launchers_map.find(static_cast(tile_N)); - FLASHINFER_CHECK( - launcher_it != launchers_map.end(), "Internal error: missing FP4 block-scale MoE launcher for tile_N=", tile_N); - auto& selected_launcher = launcher_it->second; - - // Run the launcher - it will create its own runner internally - return selected_launcher->run(config, enable_pdl); -} - -// =========================================================================== -// NVFP4 MoE LoRA (decomposed / unfused-activation) — FP4 sibling of the FP8 -// trtllm-lora op. The standard NVFP4 path fuses SwiGLU into GEMM1, which leaves -// no seam to inject the gate_up LoRA delta pre-activation. We therefore run the -// MoE as a hand-wired pipeline that mirrors what MoE::Runner::run does for the -// DeepSeek-FP8 + per-token-NvFP4 path, but with the gate_up projection executed -// as a raw (no-activation) grouped GEMM via Gemm2::Runner so the standalone, -// LoRA-aware activation kernel can run between the two GEMMs: -// -// gather (permute bf16) -> NvFP4 quant -> gate_up GEMM (K=hidden, N=2*inter, -// raw bf16 out) -> activation (adds gate_up_lora_delta pre-SwiGLU, writes -// activation_lora_input) -> NvFP4 quant -> down GEMM (K=inter, N=hidden) -> -// finalize. -// -// The hidden states are supplied as bf16 (path 3: the dispatch feeds bf16 and this op permutes -// then NvFP4-quantizes internally, with globalScaleInv = 1/448/6 + per-token scaling, matching -// SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION). No fp4-input dequant round-trip. -// =========================================================================== - -// Decomposed NvFP4 MoE-LoRA launcher. Reuses FusedMoeLauncher's routing-phase -// workspace allocation/bookkeeping (via prepare_routing-style setup) but owns -// the MoE compute pipeline. -class FP4BlockScaleLoraLauncher { - public: - // Match the plain FP4 E2m1 path's tile ladder (FP4BlockScaleLauncher::getSupportedTileNums for - // non-bf16 act). Large prefills (high avg tokens/expert) need 128/256; capping at 64 makes - // selectDefaultTileN pick a tile too small for the Gemm2 cubin to have a valid config at that - // token count -> "Failed to initialize the TMA descriptor / illegal memory access". - static constexpr std::array mBaseSupportedTileNums = {8, 16, 32, 64, 128, 256}; - - static std::vector getSupportedTileNums() { - return std::vector(mBaseSupportedTileNums.begin(), mBaseSupportedTileNums.end()); - } - - FP4BlockScaleLoraLauncher( - TensorView const& expert_indices, - TensorView const& expert_weights, - Optional const& routing_bias, - TensorView const& hidden_states, - Optional const& hidden_states_scale, - TensorView const& gemm1_weights, - TensorView const& gemm1_weights_scale, - TensorView const& gemm2_weights, - TensorView const& gemm2_weights_scale, - Optional const& output1_scales_scalar, - Optional const& output1_scales_gate_scalar, - Optional const& output2_scales_scalar, - TensorView const& gate_up_lora_delta, - TensorView const& activation_lora_input, - TensorView const& output, - int64_t lora_ready_event, - int64_t gemm2_done_event) - : expert_indices_(expert_indices), - expert_weights_(expert_weights), - routing_bias_(routing_bias), - hidden_states_(hidden_states), - hidden_states_scale_(hidden_states_scale), - gemm1_weights_(gemm1_weights), - gemm1_weights_scale_(gemm1_weights_scale), - gemm2_weights_(gemm2_weights), - gemm2_weights_scale_(gemm2_weights_scale), - output1_scales_scalar_(output1_scales_scalar), - output1_scales_gate_scalar_(output1_scales_gate_scalar), - output2_scales_scalar_(output2_scales_scalar), - gate_up_lora_delta_(gate_up_lora_delta), - activation_lora_input_(activation_lora_input), - output_(output), - lora_ready_event_(lora_ready_event), - gemm2_done_event_(gemm2_done_event) {} - - // Returns {output} when do_finalize, else {gemm2_output, expert_weights, - // expanded_idx_to_permuted_idx} for a downstream finalize kernel. - Array - run(int64_t num_experts, - int64_t top_k, - int64_t intermediate_size, - int64_t local_expert_offset, - int64_t local_num_experts, - double routed_scaling_factor, - int64_t routing_method_type, - int64_t tile_tokens_dim, - bool norm_topk_prob, - bool do_finalize, - bool enable_pdl, - bool use_fused_permute_quant) { - namespace moe_ns = tensorrt_llm::kernels::trtllmgen_moe; - auto device = hidden_states_.device(); - int dev_id = device.device_id; - cudaStream_t stream = get_stream(device); - - int64_t const num_tokens = hidden_states_.size(0); - int64_t const hidden_size = - hidden_states_.dtype() == dl_uint8 ? hidden_states_.size(1) * 2 : hidden_states_.size(1); - int64_t const inter = intermediate_size; - int64_t const gate_up_n = 2 * inter; // gated SwiGLU - - // ---- 1) routing (precomputed packed topk) ---- - Tensor num_tokens_per_expert = alloc_tensor({num_experts}, dl_int32, device); - int32_t max_num_padded_tokens = - moe_ns::Routing::getMaxPermutedPaddedCount(num_tokens, top_k, num_experts, tile_tokens_dim); - Tensor total_num_padded_tokens = alloc_tensor({1}, dl_int32, device); - Tensor expanded_idx_to_permuted_idx = alloc_tensor({num_tokens * top_k}, dl_int32, device); - Tensor permuted_idx_to_token_idx = alloc_tensor({max_num_padded_tokens}, dl_int32, device); - int64_t const hist_size = std::max(num_experts * 2, 256 * 2); - Tensor expert_count_histogram = alloc_tensor({hist_size}, dl_int32, device); - int32_t max_num_ctas = moe_ns::Routing::getMaxNumCtasInBatchDim(num_tokens, top_k, num_experts, tile_tokens_dim); - Tensor cta_idx_xy_to_batch_idx = alloc_tensor({max_num_ctas}, dl_int32, device); - Tensor cta_idx_xy_to_mn_limit = alloc_tensor({max_num_ctas}, dl_int32, device); - Tensor num_non_exiting_ctas = alloc_tensor({1}, dl_int32, device); - - auto routing_bias_dtype = routing_bias_.has_value() ? routing_bias_.value().dtype() : dl_bfloat16; - btg::Dtype mRoutingBiasDtype = routing_bias_dtype == dl_bfloat16 ? btg::Dtype::Bfloat16 : btg::Dtype::Fp32; - - // The wrapper passes an empty placeholder for expert_weights; the routing runner - // writes the unpacked per-(token,k) weights here. Allocate it ourselves (mirrors - // Fp8BlockScaleLauncher::prepare_routing when has_precomputed_weights is false). - // If the caller did pass a real expert_weights tensor, copy it into the allocation - // afterwards is unnecessary; we just compute into our own buffer for a clean Tensor - // return type. The bf16 routing-weight values are identical either way. - auto ew_dtype = mRoutingBiasDtype == btg::Dtype::Fp32 ? dl_float32 : dl_bfloat16; - Tensor expert_weights_alloc = alloc_tensor({num_tokens, top_k}, ew_dtype, device); - void* expert_weights_ptr = expert_weights_alloc.data_ptr(); - - moe_ns::Routing::Runner routing_runner(tile_tokens_dim); - routing_runner.run( - /*routing_logits=*/nullptr, - routing_bias_.has_value() ? routing_bias_.value().data_ptr() : nullptr, - num_tokens, - num_experts, - top_k, - /*n_group=*/0, - /*topk_group=*/0, - local_expert_offset, - local_num_experts, - routed_scaling_factor, - static_cast(const_cast(expert_indices_.data_ptr())), - static_cast(expert_count_histogram.data_ptr()), - static_cast(total_num_padded_tokens.data_ptr()), - static_cast(expanded_idx_to_permuted_idx.data_ptr()), - /*permuted_idx_to_expanded_idx=*/nullptr, - static_cast(permuted_idx_to_token_idx.data_ptr()), - /*expertIds=*/nullptr, - expert_weights_ptr, - static_cast(num_tokens_per_expert.data_ptr()), - static_cast(cta_idx_xy_to_batch_idx.data_ptr()), - static_cast(cta_idx_xy_to_mn_limit.data_ptr()), - static_cast(num_non_exiting_ctas.data_ptr()), - btg::Dtype::Bfloat16, - mRoutingBiasDtype, - /*useRoutingScalesOnInput=*/false, - /*useDeepSeekFp8=*/false, - static_cast(routing_method_type), - stream, - btg::Dtype::Bfloat16, - norm_topk_prob, - /*routing_replay_out=*/nullptr); - - // ---- 2) hidden as bf16 (path 3: dispatch feeds bf16; the op quantizes internally) ---- - TVM_FFI_ICHECK(hidden_states_.dtype() == dl_bfloat16) - << "fp4 LoRA (path 3) requires bf16 hidden_states; the dispatch feeds bf16 and the op " - "permutes+NvFP4-quantizes internally (no python pre-quant / dequant round-trip)."; - void* hidden_bf16_ptr = hidden_states_.data_ptr(); - - int64_t const tile = tile_tokens_dim; - // gate_up GEMM act operand (permuted fp4 hidden + scales). Declared in run() scope because the - // gate_up GEMM (step 5) consumes them; the LARGE [max_padded, hidden] bf16 gather buffer - // (permuted_hidden_bf16, ~4 GB at a 32K-token prefill) lives only inside the block below, so it - // frees right after the quant -- before the equally-large [max_padded, hidden] gemm2_output is - // allocated (step 8), letting the caching allocator reuse its block (halves the op's peak). - auto gu_sfLayout = tile >= 128 ? tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4 - : tensorrt_llm::QuantizationSFLayout::SWIZZLED_8x4; - int64_t const hidden_sf_size = tensorrt_llm::computeSwizzledLayoutSFSize(max_num_padded_tokens, hidden_size / 16); - Tensor hidden_fp4 = alloc_tensor({max_num_padded_tokens, hidden_size / 2}, dl_uint8, device); - Tensor hidden_fp4_sf = alloc_tensor({hidden_sf_size}, dl_uint8, device); - Tensor hidden_per_token_sf = alloc_tensor({max_num_padded_tokens}, dl_float32, device); - if (use_fused_permute_quant && tile < 128) { - // Invariants the fused kernel relies on (review hardening): hidden must be a multiple of the - // 16-wide PackedVec load, and top_k must fit the dedup per-token-scale write (threadIdx