Update on "Add structured stats reporting and GPU memory tracking to Qwen3.5 MoE runner"

Runner now uses llm::Stats with proper timestamps for model load, prefill,
decode, and GPU memory (via cudaMemGetInfo). Output matches stats.h
print_report format: PyTorchObserver JSON line plus human-readable table.

This commit was authored with the assistance of Claude Code.

[ghstack-poisoned]

Author: digantdesai

.ci/scripts/export_model_artifact.sh

@@ -418,7 +418,9 @@ if [ "$MODEL_NAME" = "qwen3_5_moe" ]; then
   TORCHINDUCTOR_CACHE_DIR="$INDUCTOR_CACHE" \
   python -m executorch.examples.models.qwen3_5_moe.export \
       --prequantized "$LOCAL_MODEL_DIR" \
-      --output-dir "${OUTPUT_DIR}"
+      --output-dir "${OUTPUT_DIR}" \
+      --dense-prefill dequant \
+      --moe-activation-dtype int8
   echo "::endgroup::"
 
   test -f "${OUTPUT_DIR}/model.pte"

examples/models/qwen3_5_moe/export.py

@@ -667,12 +667,12 @@ def _apply_turboquant(model, config):
 # ---------------------------------------------------------------------------
 
 
-def _set_batched_moe(model, enabled, moe_moe_moe_moe_activation_dtype="bf16"):
+def _set_batched_moe(model, enabled, moe_activation_dtype="bf16"):
     """Toggle batched tensor-core MoE kernel for all MoE layers."""
     for layer in model.layers:
         if hasattr(layer, "mlp") and hasattr(layer.mlp, "experts"):
             layer.mlp.experts.use_batched_moe = enabled
-            layer.mlp.experts.moe_moe_moe_moe_activation_dtype = moe_moe_moe_moe_activation_dtype
+            layer.mlp.experts.moe_activation_dtype = moe_activation_dtype
 
 
 def export_and_lower(model, config, args):
@@ -916,8 +916,8 @@ def _export_cuda(model, config, args):
     # chunk_gated_delta_rule with CHUNK_SIZE=64) for the full range of sequence
     # lengths. Smaller examples cause AOTI to bake in intermediate buffer sizes
     # that reject longer prompts at runtime.
-    moe_moe_moe_moe_activation_dtype = getattr(args, "moe_moe_moe_moe_activation_dtype", "bf16")
-    _set_batched_moe(model, True, moe_moe_moe_moe_activation_dtype=moe_moe_moe_moe_activation_dtype)
+    moe_activation_dtype = getattr(args, "moe_activation_dtype", "bf16")
+    _set_batched_moe(model, True, moe_activation_dtype=moe_activation_dtype)
     dense_prefill = getattr(args, "dense_prefill", "tinygemm")
     _set_dequant_prefill(model, dense_prefill == "dequant")
     print("Exporting prefill method...")
@@ -1087,14 +1087,15 @@ def main():  # noqa: C901
         "--moe-activation-dtype",
         choices=["bf16", "int8"],
         default="bf16",
-        help="MoE activation dtype for prefill only. Decode always uses bf16. bf16 (default): W4A16 batched GEMM. int8: W4A8 with INT8 tensor cores (~1.5x faster prefill).",
+        help="MoE activation dtype for prefill only. Decode always uses bf16. bf16 (default): W4A16 batched GEMM. int8: W4A8 with INT8 tensor cores.",
     )
     parser.add_argument(
         "--dense-prefill",
         choices=["tinygemm", "dequant"],
         default="tinygemm",
-        help="Dense linear kernel: tinygemm (default W4A16 INT4 kernel) or "
-        "dequant (dequant W4→BF16 + Inductor mm for prefill, int4_matvec for decode).",
+        help="Dense linear prefill kernel. Decode always uses int4_matvec (Triton W4A16 vec-mat). "
+        "tinygemm (default): W4A16 _weight_int4pack_mm. "
+        "dequant: dequant W4→BF16 + cuBLAS GEMM.",
     )
     args = parser.parse_args()
 
@@ -1139,7 +1140,7 @@ def main():  # noqa: C901
                 "(dense weights must be W4 quantized)"
             )
 
-    if args.moe_moe_moe_activation_dtype != "bf16" and args.backend != "cuda":
+    if args.moe_activation_dtype != "bf16" and args.backend != "cuda":
         parser.error("--moe-activation-dtype int8 requires --backend cuda")
 
     model, config = load_and_quantize(args)

examples/models/qwen3_5_moe/model.py

@@ -479,7 +479,7 @@ def __init__(self, config):
         self.hidden_size = config.hidden_size
         self.group_size = 32
         self.use_batched_moe = False
-        self.moe_moe_moe_activation_dtype = "bf16"
+        self.moe_activation_dtype = "bf16"
 
         self.w1_weight = nn.Parameter(
             torch.empty(
@@ -498,7 +498,7 @@ def __init__(self, config):
 
     def forward(self, x, expert_weights, expert_indices, top_k):
         if self.use_batched_moe:
-            if self.moe_moe_moe_activation_dtype == "int8":
+            if self.moe_activation_dtype == "int8":
                 return torch.ops.triton.fused_moe_batched_gemm_int8(
                     x,
                     self.w1,