forked from Karylab-cklius/vllm
cf6c0d2518
* [Model] Add MiniMax M3 text backbone skeleton + SwiGLU-OAI clamp activation Port the MiniMax M3 (text backbone) into vLLM's custom model layout: - Add MiniMaxM3SparseForCausalLM under vllm/models/minimax_m3/nvidia with the decoder/model/causal-LM wiring; attention and MoE bodies plus weight loading are left as stubs. Dense MiniMaxM3MLP is fully ported. - Add MiniMaxM3SparseForConditionalGeneration as a minimal LM-routing wrapper (KimiK25-style init_vllm_registered_model on text_config) and register both architectures. - Add MiniMaxM3Config (model_type minimax_m3_vl) wrapping MiniMaxM3TextConfig so config.get_text_config() extracts the backbone; register in the config registries. Generalize silu_and_mul_with_clamp to SwiGLU-OAI: - Add alpha (scales the activation's sigmoid) and beta (added to the non-activated half) to the CUDA kernel, ops.h, and torch_bindings schema. Defaults alpha=1.0, beta=0.0 are bitwise-identical to the previous silu(gate)*up, so existing callers (DeepSeek V4) are unaffected. - SiluAndMulWithClamp(alpha, beta) used by MiniMaxM3MLP with alpha=swiglu_alpha, beta=1.0, matching the reference gate*sigmoid(alpha*gate)*(up+1). AI assistance (Claude) was used for this change. Signed-off-by: Yongye Zhu <yongye@inferact.ai> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> * [Model] MiniMax M3: implement MoE block + weight-name mapping Implement the sigmoid-routed MoE block for MiniMax M3 and map module names to the checkpoint so weight loading works for the ported modules. MoE block (MiniMaxM3MoE): - fp32 router via GateLinear (bf16 activations upcast to fp32; fp32 weights and logits), matching minimax_m2/sglang precision. - FusedTopKBiasRouter routing (scoring_func from config, sigmoid + e_score_correction_bias + renormalize), verified to match sglang's TopK (select-with-bias, weight-without-bias, routed_scaling on output). - swigluoai activation (from config.hidden_act) + swiglu_limit; shared expert fused into FusedMoE so the shared partial is reduced with the routed output. Weight loading: - Name the MoE submodule `block_sparse_moe` (dense stays `mlp`) to match the checkpoint; decoder forward selects per layer. - load_weights handles gate_up fusion (dense MLP + shared experts) and expert w1/w2/w3 -> w13/w2 fusion; wrappers delegate via AutoWeightsLoader, skipping vision/mm/mtp. Not-yet-ported modules (attention) are skipped until they land. The expert GEMM/activation kernel correctness and attention/weight loading for the remaining modules are not part of this change. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> * [MiniMax M3] Port attention modules and wire up MXFP8 checkpoint loading Port the weight-bearing dense (MiniMaxM3Attention) and sparse (MiniMaxM3SparseAttention) attention modules so the checkpoint's self_attn.* tensors map onto real params (forward still stubbed; this targets weight loading). Add qkv stacked mapping and the weight_scale_inv -> weight_scale remap in load_weights. Load MiniMax-style MXFP8 checkpoints (quant_method: "mxfp8" + ignored_layers) via the ModelOpt MXFP8 config: register "mxfp8" in method_to_config and normalize the minimal checkpoint schema to the ModelOpt schema in ModelOptMxFp8Config.from_config (same on-disk format). Use setdefault for online shorthands so the checkpoint config wins over the "mxfp8" online shorthand. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <yongye@inferact.ai> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> * [MiniMax M3] Add DeepGEMM MXFP8 MoE backend with swigluoai support Extend DeepGemmExperts to support MXFP8 activations (FP8 e4m3 + UE8M0 1x32 block scales) via the grouped GEMM with recipe (1, 32), reusing the oracle/fp8 weight-conversion path. Generalize deepgemm_post_process_fp8_weight_block to derive the transform recipe from the block shape ((1, 1, 32) for MXFP8) and accept uint8 E8M0 scales. Unify the fused gated-activation+quant triton kernels around y = (up + beta) * gate * sigmoid(alpha * gate): silu is alpha=1, beta=0 (bit-identical to before); swigluoai uses alpha/beta from config. Thread gemm1_alpha/gemm1_beta from the FusedMoE layer through the MXFP8 quant config into the kernels, and add swiglu_alpha/swiglu_beta to the layer and MiniMax M3 config/model (beta sourced from config, not hardcoded). Wire Fp8MoeBackend.DEEPGEMM into the MXFP8 oracle (selectable via --moe-backend deep_gemm), resolving directly to DeepGemmExperts (the Triton fallback cannot handle the 1x32 scheme). Advertise SWIGLUOAI in _supports_activation so swigluoai selects DeepGEMM rather than falling through to another backend; gate the MXFP8 scheme to Blackwell (SM100). Verified on GB200: packed-kernel parity (silu defaults unchanged, swigluoai matches torch ref), (1,32) weight-prep transform, and a TP=4 launch selecting the DEEPGEMM MXFP8 backend with full weight load (the run then stops at the still-stubbed attention forward, as expected). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <yongye@inferact.ai> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> * [MiniMax M3] Implement dense attention forward Implement MiniMaxM3Attention.forward (dense path): qkv projection, split, per-head QK norm (GemmaRMSNorm, qk_norm_type="per_head"), partial RoPE, attention, and output projection. Mirrors the sglang reference dense path and vLLM's canonical per-head-norm convention. attention_output_gate is False for M3, so the gate branch is omitted. The sparse attention forward (index branch) remains stubbed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <yongye@inferact.ai> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> * [MiniMax M3] Drop dead index value/output projection branch For M3, sparse_disable_index_value matches sparse_attention_freq exactly ([0,0,0,1,...,1]): the only layers with the flag unset (0-2) are the non-sparse layers built as MiniMaxM3Attention. Every layer that constructs MiniMaxM3SparseAttention therefore always disables the index value/output projections, so index_{v,o}_proj are never created. Remove the unreachable else branch, the disable_index_value parameter and field, and the now-unused _disable_index_value_layer_ids helper. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <yongye@inferact.ai> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> * [MiniMax M3] Add sparse-attention backend + merged attention layer Scaffold the lightning-indexer sparse-attention path: - MiniMaxM3SparseBackend (registered as MINIMAX_M3_SPARSE): block-sparse GQA backend; get_kv_cache_shape serves both the main K/V cache and the single-vector index-key side cache. - MiniMaxM3IndexerCache: side KV cache for per-token index keys, key-only so it uses a single-vector MLAAttentionSpec rather than a K+V FullAttentionSpec. - MiniMaxM3SparseMetadata (+ prefill/decode sub-metadata) and its builder, splitting the batch via split_decodes_and_prefills. - MiniMaxM3SparseImpl: subclasses AttentionImplBase so it can take a custom forward(query, index_query, kv_cache, index_kv_cache); no alibi / sliding window / logits soft cap. forward is a stub pending the kernel port. MiniMaxM3SparseAttention is merged into a single AttentionLayerBase: it owns the projections, per-head QK norm and RoPE, binds the backend + impl, registers the main K/V cache, and holds the index cache. Its forward computes q/k/v and the index q/k, pre-inserts K/V and index-K into their caches, then calls the sparse impl with only the queries. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <yongye@inferact.ai> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> * [MiniMax M3] Implement sparse-attention forward + MXFP8 DeepGEMM MoE e2e Port the MiniMax M3 lightning-indexer sparse attention to a Triton backend and fix the MXFP8 DeepGEMM MoE path so the model runs end to end. Attention (vllm/v1/attention/ops/minimax_m3_sparse_ops.py + backend): - Triton kernels (paged, page == sparse block == 128): index block-score + bitonic top-k, and GQA block-sparse flash attention over the selected blocks. - MiniMaxM3SparseImpl.forward: decode-first split, dispatching the same kernels per phase (a decode token is a 1-token prefill). Index and main caches use separate block tables. - Dedicated MiniMaxM3IndexerBackend for the key-only index cache so the main GQA cache (num_kv_heads==1 at TP>=4) is not mistaken for the index layout. - get_supported_kernel_block_sizes()==[128] (one sparse block per KV page). MXFP8 DeepGEMM MoE: - Prepare-phase activation quant emits float32 per-(1,32) group scales for the DeepGEMM backend (use_deep_gemm_packed_mxfp8 on the quant config), matching the FP8 128-block path with group=32. - deepgemm_moe_permute / ep_scatter take a block_size so the activation-scale group (32) is honored through the expert permute. - workspace_shapes uses the contiguous-layout M alignment (not block_shape[0], which is 1 for MXFP8 and under-sized the workspace). GSM8K (5-shot, TP=4) flexible-extract 0.921 / strict 0.919. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <yongye@inferact.ai> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> * [MiniMax M3] Enable decode CUDA graphs + dedicated split-K decode kernels (#7) Two changes to the sparse-attention backend: 1. Full decode CUDA-graph support. The metadata builder now declares AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE and precomputes all per-step kernel args (cu_seqlens_q, prefix_lens, max_query_len, num_actual_tokens) in build(), removing the .item() host sync and the per-step torch.zeros/cumsum/diff from the impl forward. Derived decode prefix lengths are written into a persistent buffer so the captured graph reads stable addresses across replays. 2. Dedicated split-K decode kernels (mirroring the sglang reference) instead of reusing the prefill kernels with BLOCK_SIZE_Q=1, which left the GPU idle at decode (one query token per request). The index score now splits over seq blocks and the GQA attention splits over the selected top-k blocks with an LSE merge (flash-decoding). Chunk counts depend only on shape constants, so the grid is fixed within a CUDA graph. Verified: a parity test against the prior (GSM8K 92.1) prefill-as-decode path matches exactly on top-k selection and on attention output (bf16 noise) across seq lengths 128-2048. Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> * Init Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> * Move Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> * FIX Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> * Addresss conflict Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> --------- Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>