Bump sglang to v0.5.16 and add GLM-5.2 support

Makefile

@@ -9,6 +9,12 @@ fix-permissions:
 build-docker:
 	sudo docker build -t $(IMAGE_NAME) -f dockerfiles/cuda12.6.dockerfile .;
 
+# Thin image on top of the official SGLang release image. Required for Blackwell
+# (GB200/B200, sm_100+), which needs CUDA 13; also much faster to build.
+SGLANG_IMAGE ?= lmsysorg/sglang:v0.5.16-cu130
+build-docker-sglang:
+	sudo docker build -t $(IMAGE_NAME) -f dockerfiles/sglang-cu130.dockerfile --build-arg SGLANG_IMAGE=$(SGLANG_IMAGE) .;
+
 run-docker:
 	@if [ -n "$(GPU_DEVICES)" ]; then GPU_DEVICES_FLAG="--gpus=$(GPU_DEVICES)"; else GPU_DEVICES_FLAG=""; fi; \
 	echo "Running: sudo docker run $$GPU_DEVICES_FLAG --network=host --cap-add=SYS_ADMIN -it -p ${PORT}:22 --memory ${MEM_LIMIT} --name $(CONTAINER_NAME) $(IMAGE_NAME)"; \

README.md

@@ -51,6 +51,20 @@ make build-docker
 
 This creates a local image named `flowsim-image` with FlowSim patches already applied to sglang.
 
+On Blackwell (GB200/B200, sm_100+) use the CUDA 13 image instead — `cuda12.6.dockerfile`
+builds kernels for sm_90 and below:
+
+```bash
+make build-docker-sglang
+```
+
+This layers FlowSim on top of the official `lmsysorg/sglang:v0.5.16-cu130` release image
+(which already ships matching CUDA 13 / FlashInfer / DeepGEMM builds for amd64 and arm64)
+rather than building sglang from source, so it is also much faster. Override the base with
+`make build-docker-sglang SGLANG_IMAGE=...`; it must match the `workload/framework/sglang`
+submodule commit, since `patches/hook-v0516.patch` is applied to the sglang checkout inside
+the image. The build needs no network access.
+
 ### 2. Profile (Generate Traces)
 
 Use `flowsim submit` to capture stage-separated traces (EXTEND + DECODE), parse them, and run cross-rank analysis — all in one step. See [Stage Profiling](#stage-profiling) for how stages and collection modes work.
@@ -149,8 +163,32 @@ flowsim submit --scheduler local \
     --model-path workload/models/configs/Qwen3-235B-A22B \
     --sweep 1:2048:0 4:2048:0 8:2048:0 --decode-tokens 2 --gpus 1 \
     --extra-server-opts "--load-format dummy"
+
+# GLM-5.2 (MLA + DSA indexer + MoE). Its config declares a 1M context, so the
+# KV pool must be capped or the server OOMs before the first forward.
+flowsim submit --scheduler local \
+    --collect all \
+    --model-path /flowsim/workload/models/configs/GLM-5.2 \
+    --tp 1 --bs 1 --input-len 2048 --existing-ctx 0 --decode-tokens 4 --gpus 1 \
+    --extra-server-opts "--load-format dummy --mem-fraction-static 0.45 --context-length 8192"
 ```
 
+### Model configs
+
+`workload/models/configs/` holds **truncated** HF configs: layer counts are cut to the
+smallest set that still exercises every distinct layer type, while expert counts, hidden
+sizes and head counts are left at their real values. Profiling one dense + one MoE layer is
+enough to extrapolate the full model, and it keeps a run to a single GPU.
+
+For DeepSeek-style models (`deepseek`, `GLM-5.2`) that means `num_hidden_layers: 2` with
+`first_k_dense_replace: 1`, giving layer 0 dense and layer 1 MoE. sglang places MoE layers
+using `first_k_dense_replace` and `moe_layer_freq`; the HF-only `mlp_layer_types` /
+`indexer_types` lists are truncated to match purely for `transformers` self-consistency.
+Models without dense prefix layers (`Qwen3-235B-A22B`) use `num_hidden_layers: 1`.
+
+`GLM-5.2/config_full.json` is the unmodified upstream config, kept so the truncation stays
+auditable — it differs from `config.json` only in the four keys above.
+
 For K8s / Slurm clusters, replace `--scheduler local` with `k8s` or `slurm`. See [schedulers/README.md](schedulers/README.md) for full scheduler documentation.
 
 ### Output structure

dockerfiles/sglang-cu130.dockerfile

@@ -0,0 +1,64 @@
+# FlowSim on top of an upstream SGLang runtime image.
+#
+# Why a second dockerfile: dockerfiles/cuda12.6.dockerfile builds everything from
+# nvcr.io/nvidia/pytorch:24.10-py3 (CUDA 12.6), whose kernels are compiled for
+# sm_90 and below. Blackwell parts (GB200/B200, sm_100+) need CUDA 13, and
+# building SGLang + FlowInfer + DeepGEMM from source for aarch64 is slow and
+# fragile. Instead this image starts from the official SGLang release image,
+# which already ships matching CUDA 13 / FlashInfer / DeepGEMM builds for both
+# amd64 and arm64, and only layers FlowSim on top.
+#
+# The base image pins the same SGLang commit as the workload/framework/sglang
+# submodule, and installs it in editable mode from /sgl-workspace/sglang, so the
+# FlowSim tracing patch can be applied directly to that checkout.
+#
+# Build:
+#   make build-docker-sglang
+#   # or
+#   docker build -t flowsim-image -f dockerfiles/sglang-cu130.dockerfile .
+ARG SGLANG_IMAGE=lmsysorg/sglang:v0.5.16-cu130
+FROM ${SGLANG_IMAGE}
+
+LABEL maintainer="FlowSim"
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# No apt/pip install step on purpose. GPU hosts frequently build without egress,
+# and everything needed is already in the base image:
+#   * git + patch          -> applying the FlowSim tracing patch
+#   * requests/numpy/pandas/PyYAML -> the profile + trace-parsing code paths
+# `perfetto` is listed in pyproject.toml but is not imported by scripts/ or
+# simulator/, so it is not required to produce the kernel CSVs. Pass extra
+# packages when the builder does have egress, e.g.
+#   docker build --build-arg PIP_EXTRA_INSTALL="perfetto scalesim scipy" ...
+ARG PIP_EXTRA_INSTALL=""
+RUN if [ -n "${PIP_EXTRA_INSTALL}" ]; then \
+        pip install --no-cache-dir ${PIP_EXTRA_INSTALL}; \
+    fi && \
+    command -v git >/dev/null && command -v patch >/dev/null && \
+    python -c "import requests, numpy, pandas, yaml; print('FlowSim python deps OK')"
+
+# Copy FlowSim itself. The sglang submodule is deliberately not copied: this
+# image uses the SGLang checkout that ships in the base image.
+COPY scripts /flowsim/scripts
+COPY simulator /flowsim/simulator
+COPY schedulers /flowsim/schedulers
+COPY utils /flowsim/utils
+COPY backend /flowsim/backend
+COPY tests /flowsim/tests
+COPY workload/models /flowsim/workload/models
+COPY workload/framework/patches /flowsim/workload/framework/patches
+COPY kernels.json pyproject.toml README.md /flowsim/
+
+# Apply the FlowSim kernel-tracing hooks to the SGLang install in the base
+# image. Fails the build if the base image drifts from the pinned submodule.
+ARG SGLANG_SRC=/sgl-workspace/sglang
+RUN cd ${SGLANG_SRC} && \
+    git apply --verbose /flowsim/workload/framework/patches/hook-v0516.patch && \
+    git apply --verbose /flowsim/workload/framework/patches/0006-Balanced-moe-for-glm4-moe-v0516.patch && \
+    python -c "from sglang.srt.tracing.hook_register import register_kernels_for_profiling; print('FlowSim tracing hooks installed')" && \
+    python -c "from sglang.srt.layers.moe.topk import balance_router_logits; print('FlowSim balanced-MoE installed')"
+
+WORKDIR /flowsim
+ENV PYTHONPATH=/flowsim
+CMD ["/bin/bash"]

tests/integration/test_model_config.py

@@ -14,6 +14,7 @@
     [
         "/flowsim/workload/models/configs/deepseek",
         "/flowsim/workload/models/configs/gpt3",
+        "/flowsim/workload/models/configs/GLM-5.2",
     ],
 )
 def test_docker_image(tp, model_path):
@@ -28,12 +29,16 @@ def test_docker_image(tp, model_path):
         "dummy",
         "--tp",
         str(tp),
-        "--batch",
+        "--batch-size",
         "1",
         "--input-len",
         "128",
         "--output-len",
         "2",
+        # GLM-5.2 declares a 1M max_position_embeddings; without a cap the KV
+        # pool is sized for it and the run OOMs before the first forward.
+        "--context-length",
+        "8192",
     ]
     result = subprocess.run(cmd, capture_output=True, text=True)
     output = result.stdout + result.stderr

tests/unit/test_shape_merge.py

@@ -221,6 +221,52 @@ def test_unmatched_kernels(self, tmp_path):
         rows = _read_csv(output_csv)
         assert rows[0]["Dims"] == "N/A"
 
+    def _write_count_mismatch_pair(self, tmp_path):
+        """Timing pass sees `fused` once; eager pass expands it into 2 `step`s."""
+        timing_csv = str(tmp_path / "timing.csv")
+        shape_csv = str(tmp_path / "shape.csv")
+        _write_csv(
+            timing_csv,
+            [
+                {"Name": "step", "Duration (us)": "100", "Dims": "N/A"},
+                {"Name": "gemm", "Duration (us)": "300", "Dims": "N/A"},
+            ],
+        )
+        _write_csv(
+            shape_csv,
+            [
+                {"Name": "step", "Duration (us)": "99", "Dims": "[1]"},
+                {"Name": "step", "Duration (us)": "98", "Dims": "[2]"},
+                {"Name": "gemm", "Duration (us)": "299", "Dims": "[64, 128]"},
+            ],
+        )
+        return timing_csv, shape_csv
+
+    def test_count_mismatch_skipped_by_default(self, tmp_path):
+        """A kernel launched a different number of times is left untouched.
+
+        Disabling CUDA graphs can change the launch sequence (e.g. SGLang's DSA
+        decode backend replays one fused metadata kernel under CUDA graph but
+        issues several eager ops without it).  Positional matching would then
+        attach the wrong shapes, so those kernels keep their timing-pass values
+        while unaffected kernels still merge.
+        """
+        timing_csv, shape_csv = self._write_count_mismatch_pair(tmp_path)
+        output_csv = str(tmp_path / "merged.csv")
+
+        merge_shapes(timing_csv, shape_csv, output_csv)
+        rows = _read_csv(output_csv)
+        assert rows[0]["Dims"] == "N/A"
+        assert rows[1]["Dims"] == "[64, 128]"
+
+    def test_count_mismatch_raises_when_strict(self, tmp_path):
+        """strict=True restores the fail-fast behaviour."""
+        timing_csv, shape_csv = self._write_count_mismatch_pair(tmp_path)
+        output_csv = str(tmp_path / "merged.csv")
+
+        with pytest.raises(ValueError, match="different batch counts"):
+            merge_shapes(timing_csv, shape_csv, output_csv, strict=True)
+
     def test_default_output_path(self, tmp_path):
         """When output_csv is None, default naming is used."""
         timing_csv = str(tmp_path / "timing.csv")