[Perf] minimax_m31: Run o_proj in the launch, route on the norm's scale, read B operands from LDS

README.md

@@ -18,11 +18,11 @@ The library implements SGLang's monokernel interface ([RFC #43949](https://githu
 | GLM-5.3-Flash | MI355X (gfx950) | 2, 4, 8 | All 45 layers (one launch each; layer 0 from the tokens, with the embedding) and the final norm, LM head and logits (one launch), in decode steps of up to 8 rows (4 at TP2) |
 | GLM-5.3 | MI355X (gfx950) | 4, 8 | All 75 MoE layers (one launch each), in decode and MTP verify steps of up to 8 rows (4 at TP4) |
 | MiniMax-M3 (MXFP4) | MI350X, MI355X (gfx950) | 4 | The 57 sparse MoE layers in decode and EAGLE3 verify steps of up to 16 rows |
-| MiniMax-M3.1-Flash (NVFP4 experts) | MI350X (gfx950) | 4 | The FFN of the 57 sparse layers, after SGLang's attention, in decode and DSpark verify steps of up to 16 rows |
+| MiniMax-M3.1-Flash (NVFP4 experts) | MI350X (gfx950) | 4 | o_proj and the FFN of the 57 sparse layers, after SGLang's attention, in decode and DSpark verify steps of up to 16 rows |
 
 ### Design
 
-- **One or two launches per layer.** DeepSeek-V4.1-Flash runs two. K1 runs the attention's input mixing and projections and writes the KV cache. K2 runs the sparse attention, the output projections, the FFN mixing, the MoE and both all-reduces. On the indexer layers, K2i runs between them: it scores the FP4 index keys in one launch, and SGLang's top-k kernels select each row's top-k from its logits. GLM-5.3-Flash and GLM-5.3 run each covered layer as one launch; GLM-5.3-Flash also runs the final norm, the LM head and the logits all-gather as one tail launch. MiniMax-M3 runs one too: K1 (the input norm, the fused qkv and index projection, the q / k norms, RoPE, the KV and index-key writes and the index scores) runs in front of the block selection, the sparse attention, o_proj, the router, the MoE and both all-reduces, in one launch. MiniMax-M3.1-Flash keeps SGLang's attention and runs the rest of each sparse layer as one launch: the attention all-reduce and residual add, the post-attention norm, the router, the shared and routed experts and the MoE all-reduce.
+- **One or two launches per layer.** DeepSeek-V4.1-Flash runs two. K1 runs the attention's input mixing and projections and writes the KV cache. K2 runs the sparse attention, the output projections, the FFN mixing, the MoE and both all-reduces. On the indexer layers, K2i runs between them: it scores the FP4 index keys in one launch, and SGLang's top-k kernels select each row's top-k from its logits. GLM-5.3-Flash and GLM-5.3 run each covered layer as one launch; GLM-5.3-Flash also runs the final norm, the LM head and the logits all-gather as one tail launch. MiniMax-M3 runs one too: K1 (the input norm, the fused qkv and index projection, the q / k norms, RoPE, the KV and index-key writes and the index scores) runs in front of the block selection, the sparse attention, o_proj, the router, the MoE and both all-reduces, in one launch. MiniMax-M3.1-Flash keeps SGLang's attention and runs the rest of each sparse layer as one launch: o_proj, the attention all-reduce and residual add, the post-attention norm, the router, the shared and routed experts and the MoE all-reduce.
 - **Graph-safe routing.** Each step is routed by its CUDA graph key alone, so every rank makes the same decision and a captured graph replays the decision made at capture.
 - **Fail-stop synchronization.** Every wait on another workgroup or GPU has a deadline. A failed wait stops the launch on every rank, so no token from a failed step is returned.
 - **In-place weights.** The kernels read the weights SGLang loaded. The few weights that need another layout or precision (MiniMax-M3's attention projections run in FP8; MiniMax-M3.1-Flash's shared expert keeps the checkpoint's MXFP8 values) are copied once at startup, and each copy is logged.

sgl_amd_mono/minimax_m31/__init__.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 """MiniMax-M3.1-Flash decode FFN monokernel on gfx950: one persistent launch per sparse
-layer from the unreduced attention o_proj partial to the reduced MoE output (attention
-TP all-reduce and residual add, Gemma post-attention norm, FP32 router and top-4, the
+layer from the attention output to the reduced MoE output (o_proj, attention TP
+all-reduce and residual add, Gemma post-attention norm, FP32 router and top-4, the
 shared and routed A16W4 experts, MoE TP all-reduce).
 
 The kernel comes from vLLM (Apache-2.0, Copyright contributors to the vLLM project) at

sgl_amd_mono/minimax_m31/a16w4.py

@@ -23,15 +23,16 @@
 STEP_DWORDS = 256  # a 16-row group's dwords per step
 
 
-def weight_loads(w, row_group, k, st0, steps, cm=CM_NT):
+def weight_loads(w, row_group, k, st0, steps, cm=CM_NT, nbytes=None):
     """Lane's 16-byte weight words of 16-row group ``row_group`` (counted over
-    the flattened ``[E N]`` rows) for steps ``st0 ..``; loads only."""
+    the flattened ``[E N]`` rows) for steps ``st0 ..``; loads only. ``nbytes``:
+    the buffer's size (0: zeros, nothing read)."""
     lane = fx.thread_idx.x % 64
     base = row_group * (2 * k) + lane * 4
     return [
         fx.Vector(
             bo.buffer_load(
-                rsrc(w),
+                rsrc(w, nbytes),
                 base + (st0 + i) * STEP_DWORDS,
                 vec_width=4,
                 dtype=T.i32,
@@ -42,15 +43,15 @@ def weight_loads(w, row_group, k, st0, steps, cm=CM_NT):
     ]
 
 
-def scale_loads(ws, row_group, k, st0, steps, cm=CM_NT):
+def scale_loads(ws, row_group, k, st0, steps, cm=CM_NT, nbytes=None):
     """Lane's scale dwords for steps ``st0 .. st0 + steps`` (st0, steps even)."""
     lane = fx.thread_idx.x % 64
     c8 = k // 256
     base = row_group // 2 * c8 * 64 + lane
     return [
         fx.Int32(
             bo.buffer_load(
-                rsrc(ws),
+                rsrc(ws, nbytes),
                 base + (st0 // 2 + i) * 64,
                 vec_width=1,
                 dtype=T.i32,