recipes/docker-images/mimo26-spark/patches/sglang/mimo26-fp8-oproj.diff

diff --git a/python/sglang/srt/layers/quantization/mimo26_fp8_serve.py b/python/sglang/srt/layers/quantization/mimo26_fp8_serve.py
new file mode 100644
index 0000000000..2aa32bfa98
--- /dev/null
+++ b/python/sglang/srt/layers/quantization/mimo26_fp8_serve.py
@@ -0,0 +1,50 @@
+"""Serve checkpoint-BF16 linear layers as 128x128 block FP8 (MiMo-V2.6 on SM121).
+
+Only the module quant_method is swapped, so the parallel linear keeps its own
+forward: input splitting, bias handling and any in-layer all-reduce are unchanged.
+Scales are fp32 amax/448 per block; the GEMM is the same Triton block-FP8 kernel
+that serves the checkpoint FP8 layers.
+"""
+import logging
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+class _BlockFp8Apply:
+    def __init__(self, base):
+        self.base = base
+
+    def apply(self, layer, x, bias=None):
+        from sglang.srt.layers.quantization.fp8_utils import triton_w8a8_block_fp8_linear
+
+        out = triton_w8a8_block_fp8_linear(x, layer._mimo26_fp8_weight, [128, 128], layer._mimo26_fp8_scale)
+        return out if bias is None else out + bias
+
+    def __getattr__(self, name):
+        return getattr(self.base, name)
+
+
+def serve_linear_as_block_fp8(mod) -> bool:
+    w = getattr(mod, "weight", None)
+    if w is None or w.dtype != torch.bfloat16 or w.dim() != 2 or not hasattr(mod, "quant_method"):
+        return False
+    if type(mod.quant_method).__name__ != "UnquantizedLinearMethod":
+        return False
+    out_f, in_f = w.shape
+    if out_f % 128 or in_f % 128:
+        return False
+    blocks = w.data.float().view(out_f // 128, 128, in_f // 128, 128)
+    scale = (blocks.abs().amax(dim=(1, 3), keepdim=True) / 448.0).clamp(min=1e-12)
+    q = (blocks / scale).to(torch.float8_e4m3fn).view(out_f, in_f).contiguous()
+    mod.register_buffer("_mimo26_fp8_weight", q, persistent=False)
+    mod.register_buffer("_mimo26_fp8_scale", scale.view(out_f // 128, in_f // 128).contiguous(), persistent=False)
+    mod.quant_method = _BlockFp8Apply(mod.quant_method)
+    return True
+
+
+def serve_modules_as_block_fp8(model, predicate, label):
+    n = sum(serve_linear_as_block_fp8(m) for name, m in model.named_modules() if predicate(name))
+    logger.info("mimo26: serving %d %s linear layers as block FP8", n, label)
+    return n
diff --git a/python/sglang/srt/models/dflash.py b/python/sglang/srt/models/dflash.py
index 5722b43d45..cd5c6cc6ae 100644
--- a/python/sglang/srt/models/dflash.py
+++ b/python/sglang/srt/models/dflash.py
@@ -880,6 +880,18 @@ class DFlashDraftModel(nn.Module):
                     f"{sorted(missing)}."
                 )
 
+        import os as _os
+
+        if _os.environ.get("MIMO26_FP8_DRAFT", "0") == "1":
+            from sglang.srt.layers.quantization.mimo26_fp8_serve import serve_modules_as_block_fp8
+
+            serve_modules_as_block_fp8(
+                self,
+                lambda n: n.startswith("layers.")
+                and n.rsplit(".", 1)[-1] in ("o_proj", "gate_up_proj", "down_proj"),
+                "DFlash draft",
+            )
+
 
 class DFlashLagunaAttention(DFlashAttention):
     """Laguna DFlash attention with the trained Laguna softplus gate."""
diff --git a/python/sglang/srt/models/mimo_v2.py b/python/sglang/srt/models/mimo_v2.py
index 6f67bef0d5..bb48d20564 100644
--- a/python/sglang/srt/models/mimo_v2.py
+++ b/python/sglang/srt/models/mimo_v2.py
@@ -13,6 +13,7 @@
 # ==============================================================================
 
 import logging
+import os
 import math
 import re
 from array import array
@@ -1708,6 +1709,14 @@ class MiMoV2ForCausalLM(nn.Module, AudioEncoderMixin):
                 expected_fused_tp_size,
                 config=self.config,
             )
+        if os.environ.get("MIMO26_FP8_OPROJ", "0") == "1":
+            from sglang.srt.layers.quantization.mimo26_fp8_serve import serve_modules_as_block_fp8
+
+            serve_modules_as_block_fp8(
+                self,
+                lambda n: n.startswith("model.layers.") and n.endswith("self_attn.o_proj"),
+                "attention o_proj",
+            )
 
     def get_embed_and_head(self):
         assert self.model is not None and self.lm_head is not None, (