bench: MLA RoPE fused vs vanilla sweep

Adds cases/mla_rope.py (DiffCase) and registers it in run_cases.py with
mla_rope-specific correctness configs and dummy `dim=[0]` (MLA head dims
are fixed by motif3: H_q=80, H_kv=16, D_nope=128, D_rope=64, D_v=128).

Also adds bench_mla_rope.yaml (single-GPU TrainJob) and a small mkdir
in run_cases.py so triton's perf_report can savefig into a fresh path.

B200, motif3 config (bs=8, sl=4096): fwd 6.90x, bwd 5.53x.

Run: `python run_cases.py --case mla_rope --dtype bf16 --save-path ...`

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

benchmarks/bench_mla_rope.yaml

@@ -0,0 +1,81 @@
+apiVersion: trainer.kubeflow.org/v1alpha1
+kind: TrainJob
+metadata:
+  name: jangwoong-mla-rope-bench
+  namespace: kbm-g-np-motif
+spec:
+  managedBy: trainer.kubeflow.org/trainjob-controller
+  podTemplateOverrides:
+    - spec:
+        containers:
+          - name: node
+            volumeMounts:
+              - mountPath: /dev/shm
+                name: shm
+              - mountPath: /mair
+                name: mair
+        volumes:
+          - emptyDir:
+              medium: Memory
+              sizeLimit: 64Gi
+            name: shm
+          - name: mair
+            persistentVolumeClaim:
+              claimName: mair
+      targetJobs:
+        - name: node
+  runtimeRef:
+    apiGroup: trainer.kubeflow.org
+    kind: ClusterTrainingRuntime
+    name: torch-distributed
+  suspend: false
+  trainer:
+    args:
+      - /bin/bash
+      - '-c'
+      - |
+        set -e
+        ACTIVATIONPATH=/mair/team-sys/jangwoong/activation
+        DATESTAMP=$(date +'%y_%m_%d_%H_%M')
+        SAVE_PATH=$ACTIVATIONPATH/benchmarks/results/mla_rope/${DATESTAMP}
+        mkdir -p $SAVE_PATH
+
+        pip install triton pandas
+
+        # Build activation from local source (copy to /tmp to avoid NFS race)
+        mkdir -p /tmp/activation_src && rm -rf /tmp/activation_src/* && \
+          rsync -a --exclude=build $ACTIVATIONPATH/ /tmp/activation_src/ && \
+          pip install --no-build-isolation /tmp/activation_src 2>&1 | tail -50
+
+        python -c "import activation; print('fused_q_rope_inplace:', activation.fused_q_rope_inplace); print('fused_kv_split_rope_cat:', activation.fused_kv_split_rope_cat)"
+
+        nvidia-smi | tee $SAVE_PATH/nvidia_smi.txt
+
+        echo "=== MLA RoPE benchmark ==="
+        cd $ACTIVATIONPATH/benchmarks
+        CUDA_VISIBLE_DEVICES=0 python run_cases.py --case mla_rope --dtype bf16 \
+          --save-path $SAVE_PATH/bench 2>&1 | tee $SAVE_PATH/bench.log
+
+        echo "=== Done. Results at: $SAVE_PATH ==="
+        exit 0
+    env:
+      - name: PYTHONUNBUFFERED
+        value: '1'
+      - name: PYTORCH_ALLOC_CONF
+        value: expandable_segments:True
+      - name: CUDA_LAUNCH_BLOCKING
+        value: '0'
+      - name: OMP_NUM_THREADS
+        value: '1'
+    image: ghcr.io/motiftechnologies/llm-training:v0.1.8
+    numNodes: 1
+    numProcPerNode: 1
+    resourcesPerNode:
+      limits:
+        cpu: '96'
+        memory: 1024Gi
+        nvidia.com/gpu: '8'
+      requests:
+        cpu: '96'
+        memory: 1024Gi
+        nvidia.com/gpu: '8'

benchmarks/cases/mla_rope.py

@@ -0,0 +1,97 @@
+"""MLA RoPE case: fused (activation Triton) vs vanilla (PyTorch native).
+
+MLA head dims are fixed by motif3 spec (H_q=80, H_kv=16, D_nope=128,
+D_rope=64, D_v=128); the benchmark's only sweep axes are (bs, sl).
+The framework's ``dim`` axis is a dummy here — pass 0 in configs.
+"""
+
+import torch
+from torch import nn
+
+import activation
+
+from common.diff_engine import DiffCase
+
+# ---- MLA shapes (motif3_seq) -----------------------------------------------
+H_Q, H_KV = 80, 16
+D_NOPE, D_ROPE, D_V = 128, 64, 128
+D_QK = D_NOPE + D_ROPE  # 192
+
+
+# ---- reference (PyTorch native) --------------------------------------------
+def _precompute_freqs_cis(dim, end, theta=10000.0):
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: dim // 2].float() / dim))
+    t = torch.arange(end, dtype=torch.float32)
+    freqs = torch.outer(t, freqs)
+    return torch.polar(torch.ones_like(freqs), freqs)
+
+
+def _apply_rotary_emb_single(x, freqs_cis):
+    x_ = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+    freqs_cis = freqs_cis[: x_.shape[1]].view(1, x_.shape[1], 1, x_.shape[3])
+    out = torch.view_as_real(x_ * freqs_cis).flatten(3)
+    return out.type_as(x)
+
+
+def _reorder(qk, rope_dim):
+    B, S = qk.shape[0], qk.shape[1]
+    qk = qk.view(B, S, -1, rope_dim // 2, 2).transpose(3, 4)
+    return qk.reshape(B, S, -1, rope_dim)
+
+
+def _vanilla(q, kv_latent, k_pe, freqs_cis):
+    q_nope, q_pe = torch.split(q, [D_NOPE, D_ROPE], dim=-1)
+    q_pe = _reorder(_apply_rotary_emb_single(q_pe, freqs_cis), D_ROPE)
+    q_total = torch.cat([q_nope, q_pe], dim=-1)
+    k_pe_roped = _reorder(_apply_rotary_emb_single(k_pe.unsqueeze(2), freqs_cis), D_ROPE)
+    k_nope, v = torch.split(kv_latent, [D_NOPE, D_V], dim=-1)
+    k_full = torch.cat([k_nope, k_pe_roped.expand(-1, -1, H_KV, -1)], dim=-1)
+    return q_total, k_full, v
+
+
+def _fused(q, kv_latent, k_pe, freqs_cis):
+    q_total = activation.fused_q_rope_inplace(q, freqs_cis, D_NOPE, D_ROPE)
+    # k_pe RoPE stays PyTorch native (head-shared, too small for custom kernel)
+    k_pe_roped = _reorder(_apply_rotary_emb_single(k_pe.unsqueeze(2), freqs_cis), D_ROPE)
+    k_full, v = activation.fused_kv_split_rope_cat(kv_latent, k_pe_roped, D_NOPE, D_V, D_ROPE)
+    return q_total, k_full, v
+
+
+class _VanillaModule(nn.Module):
+    def forward(self, q, kv_latent, k_pe, freqs_cis):
+        return _vanilla(q, kv_latent, k_pe, freqs_cis)
+
+
+class _FusedModule(nn.Module):
+    def forward(self, q, kv_latent, k_pe, freqs_cis):
+        return _fused(q, kv_latent, k_pe, freqs_cis)
+
+
+class MLARoPE(DiffCase):
+    # Framework calls build_inputs(bs, sl, dim, dtype, eps) — dim unused.
+    def build_inputs(self, bs, sl, hidden, dtype, eps):
+        return {
+            "q":         (torch.randn(bs, sl, H_Q, D_QK, dtype=dtype) * 0.5).requires_grad_(True),
+            "kv_latent": (torch.randn(bs, sl, H_KV, D_NOPE + D_V, dtype=dtype) * 0.5).requires_grad_(True),
+            "k_pe":      (torch.randn(bs, sl, D_ROPE, dtype=dtype) * 0.5).requires_grad_(True),
+            "freqs_cis": _precompute_freqs_cis(D_ROPE, sl),
+        }
+
+    def make_naive(self, I):
+        return _VanillaModule()
+
+    def make_cuda(self, I):
+        return _FusedModule()
+
+    def forward(self, obj, I):
+        # fused_q_rope_inplace needs non-leaf q; wrap both paths for fairness
+        q_in = I["q"] * 1.0
+        kv_in = I["kv_latent"] * 1.0
+        kpe_in = I["k_pe"] * 1.0
+        return obj(q_in, kv_in, kpe_in, I["freqs_cis"])
+
+    def grad_inputs(self, I):
+        return [I["q"], I["kv_latent"], I["k_pe"]]
+
+
+CASE = MLARoPE()

benchmarks/run_cases.py

@@ -62,7 +62,8 @@ def main():
     ap = argparse.ArgumentParser()
     ap.add_argument(
         "--case",
-        choices=["rms", "add_rms", "poly", "mul_poly", "grouped_mul_poly"],
+        choices=["rms", "add_rms", "poly", "mul_poly", "grouped_mul_poly",
+                 "mla_rope"],
         required=True)
     ap.add_argument("--plot", action="store_true")
     ap.add_argument(
@@ -95,12 +96,26 @@ def main():
     case: DiffCase = mod.CASE
 
     # Correctness checks across multiple configs
-    for bs, sl, hid in [(2, 128, 4096), (8, 4096, 1280), (1, 32768, 1280)]:
+    # NOTE: calculate_diff positionally calls build_inputs(hidden_size, bs, sl);
+    # bench framework positionally calls build_inputs(bs, sl, dim).  These
+    # disagree — rms-style cases don't care (all 3 axes are flat dims), but
+    # mla_rope does.  We match the bench convention in cases/mla_rope.py, so
+    # we swap arg names at the correctness call site below for that case.
+    if args.case == "mla_rope":
+        cfgs = [(1, 1024, 0), (4, 4096, 0), (8, 4096, 0)]  # (bs, sl, dummy)
+    else:
+        cfgs = [(2, 128, 4096), (8, 4096, 1280), (1, 32768, 1280)]
+    for bs, sl, hid in cfgs:
         print(
             f"Checking correctness: bs={bs}, sl={sl}, D={hid} "
             f"(N={bs*sl})...",
             end=" ")
-        calculate_diff(case, batch_size=bs, seq_len=sl, hidden_size=hid)
+        if args.case == "mla_rope":
+            # Swap so positional (hidden_size, batch_size, seq_len) maps to
+            # our build_inputs(bs, sl, dim) as (bs, sl, dummy).
+            calculate_diff(case, batch_size=sl, seq_len=hid, hidden_size=bs)
+        else:
+            calculate_diff(case, batch_size=bs, seq_len=sl, hidden_size=hid)
         print("✅")
 
     for dtype_name, dtype in dtypes:
@@ -109,6 +124,7 @@ def main():
         print(f"{'=' * 60}\n")
 
         save_dir = os.path.join(args.save_path, args.case, dtype_name)
+        os.makedirs(save_dir, exist_ok=True)
         is_grouped = args.case == "grouped_mul_poly"
 
         if args.plot:
@@ -118,6 +134,8 @@ def main():
                 dim = [1280]
             elif "poly" in args.case:
                 dim = [8192, 16384]
+            elif args.case == "mla_rope":
+                dim = [0]  # MLA head dims are fixed; dim axis is a dummy
             else:
                 dim = [2048, 4096]
             configs = list(
@@ -170,6 +188,8 @@ def main():
                 dim = [1280]
             elif "poly" in args.case:
                 dim = [8192, 16384]
+            elif args.case == "mla_rope":
+                dim = [0]  # MLA head dims are fixed; dim axis is a dummy
             else:
                 dim = [2048, 4096]
             configs = list(