Laguna-S-2.1-NVFP4 on DGX Spark (GB10): 2-node TP=2 vLLM at native 1M context over ConnectX-7 RoCE (ray backend), DFlash spec decoding + FlashInfer B12X NVFP4 MoE, fronted by llm-d EPP. Includes GB10 UMA wedge-safety rationale (why ray not mp, no VLLM_SKIP_INIT_MEMORY_CHECK).

laguna-s-2.1-vllm-tp2-1M.yaml

# Laguna-S-2.1-NVFP4 on NVIDIA DGX Spark (GB10, sm_121a, arm64)
# 2-node TENSOR-PARALLEL (TP=2) vLLM serving at NATIVE 1M context, one engine
# sharded across two GB10 nodes over ConnectX-7 RoCE, with DFlash speculative
# decoding and the FlashInfer B12X NVFP4 MoE backend. Fronted by an llm-d
# Endpoint Picker (EPP) constrained to rank0.
#
# ─────────────────────────────────────────────────────────────────────────────
# WHY TP=2 (vs 2 independent replicas):
#   One vLLM engine tensor-sharded across 2 DGX Spark nodes (1 GPU/node) gives a
#   single UNIFIED KV pool and shards the weights (~35 GiB/node vs ~70 GiB on a
#   single node), freeing memory for KV. At 1M context this yields materially
#   higher concurrency than two independent TP=1 replicas, and lets one request
#   use the combined KV pool.
#
# WHY RAY (NOT the `mp` distributed backend):
#   On this exact GB10 node pair, `--distributed-executor-backend mp` WEDGED the
#   nodes at both 262144 and 131072 context (UMA memory exhaustion during model
#   init starves the kubelet -> node NotReady -> physical power-cycle). The `mp`
#   configs that appear to "work" elsewhere rely on VLLM_SKIP_INIT_MEMORY_CHECK=1,
#   which we DELIBERATELY DO NOT SET (see below). `ray` is the only multi-node
#   backend that boots cleanly here WITHOUT that skip flag.
#     - rank0 (ordinal 0) = Ray head + vLLM API leader.
#     - rank1 (ordinal 1) = `ray start --block` worker (no vLLM process, no API).
#       RayDistributedExecutor schedules the 2 TP workers across the Ray cluster.
#
# UMA WEDGE-SAFETY (GB10 shares ONE ~121 GiB pool between system RAM and GPU):
#   - NO VLLM_SKIP_INIT_MEMORY_CHECK: keep vLLM's pre-flight guard so an
#     over-allocation is a CLEAN pod-fail, not a kubelet-starving node wedge.
#   - --enforce-eager: skip CUDA-graph capture (its transient buffer spike was a
#     wedge trigger on UMA).
#   - --kv-cache-memory pinned BELOW the "fully utilize" ceiling to leave headroom
#     for allocation variance/fragmentation (do not pin KV to the exact max).
#   - PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True reduces fragmentation
#     overshoot during the memory-profiling forward pass.
#   - The b12x MoE workspace is shared across layers (vLLM PR #48698) so it no
#     longer blows up per-layer during profile_run — the original wedge cause.
#
# NATIVE 1M CONTEXT:
#   The published NVFP4 weights are native 1M checkpoints (long-context extension
#   trained + quantization calibrated at 1,048,576 per the model card), but
#   config.json ships the 256K profile by default. The fetcher Job patches
#   config.json to the 1M profile post-download (idempotent, survives re-fetch):
#     target: rope full_attention factor 32->128, attention_factor updated,
#             max_position_embeddings 262144->1048576
#     draft (DFlash head, rope=null): max_position_embeddings only.
#   This RESTORES the trained config; it is NOT extrapolation past training.
#
# Placeholders to replace for your environment:
#   <REGISTRY>   your OCI registry host (image built from a vLLM source fork with
#                sm_121a kernels + the b12x workspace-sharing patch)
#   <IMAGE_TAG>  your built image tag
#   <NAMESPACE>  target namespace
#   <NODE_A>,<NODE_B>  the two GB10 hostnames (one GPU each, cross-linked by RoCE)
#   <ROCE_NAD>   Multus NetworkAttachmentDefinition for the RoCE macvlan between
#                the two nodes (mtu 9000). NCCL_IB_HCA below must name the RoCE
#                devices actually present in the pod netns.
#   <HF_TOKEN_SECRET>  Secret holding a HuggingFace token (key HF_TOKEN)
# ─────────────────────────────────────────────────────────────────────────────
---
apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: laguna-model-cache
  namespace: <NAMESPACE>
spec:
  # RWX so the fetcher writes once and both ranks read (shared filesystem).
  accessModes: ["ReadWriteMany"]
  storageClassName: <RWX_STORAGE_CLASS>
  resources:
    requests:
      storage: 100Gi
---
# One-shot fetcher: downloads both checkpoints into the shared PVC and patches
# config.json to the native 1M profile. Job pod templates are immutable, so on
# any change delete+recreate the Job (or use your GitOps tool's Replace option).
apiVersion: batch/v1
kind: Job
metadata:
  name: laguna-model-fetcher
  namespace: <NAMESPACE>
spec:
  ttlSecondsAfterFinished: 86400
  backoffLimit: 20
  manualSelector: true
  selector:
    matchLabels:
      job: laguna-model-fetcher
  template:
    metadata:
      labels:
        job: laguna-model-fetcher
    spec:
      restartPolicy: OnFailure
      nodeSelector:
        kubernetes.io/arch: arm64
      containers:
        - name: fetcher
          image: <REGISTRY>/vllm-spark-runtime:<IMAGE_TAG>
          securityContext:
            runAsUser: 0
          command: [/usr/bin/python3]
          args:
            - -c
            - |
              import os
              import json
              import glob
              import socket
              from huggingface_hub import snapshot_download

              # Force HuggingFace downloads over IPv4 (dual-stack egress workaround).
              _getaddrinfo = socket.getaddrinfo
              socket.getaddrinfo = lambda *a, **k: [x for x in _getaddrinfo(*a, **k) if x[0] == socket.AF_INET]
              for model in (os.environ["TARGET_MODEL"], os.environ["DRAFT_MODEL"]):
                  snapshot = snapshot_download(model, max_workers=8)
                  print(f"snapshot={snapshot}")

              # Restore the model's NATIVE 1M-context configuration (per the model
              # card). The NVFP4 weights are native 1M checkpoints; config.json
              # ships the 256K profile by default. Editing rope factor 32->128
              # (+ attention_factor) and max_position_embeddings 262144->1048576
              # restores the TRAINED 1M config -- NOT extrapolation. Idempotent;
              # patched post-download so it survives every re-fetch.
              def patch_config(model, rope):
                  cfg = sorted(glob.glob(f"/models/huggingface/hub/models--{model.replace('/','--')}/snapshots/*/config.json"))
                  if not cfg:
                      raise SystemExit(f"config.json not found for {model}")
                  path = cfg[-1]
                  c = json.load(open(path))
                  c["max_position_embeddings"] = 1048576
                  if rope:
                      fa = c.setdefault("rope_parameters", {}).setdefault("full_attention", {})
                      fa["factor"] = 128.0
                      fa["attention_factor"] = 1.4852030263919618
                  json.dump(c, open(path, "w"), indent=2)
                  print(f"patched 1M config: {path} (rope={'yes' if rope else 'max_pos_only'})")

              # Target has YaRN full_attention rope; draft (DFlash head) has
              # rope_parameters=null -> only raise max_position_embeddings.
              patch_config(os.environ["TARGET_MODEL"], rope=True)
              patch_config(os.environ["DRAFT_MODEL"], rope=False)
              open("/models/.laguna_fetched", "w").close()
          env:
            - name: TARGET_MODEL
              value: poolside/Laguna-S-2.1-NVFP4
            - name: DRAFT_MODEL
              value: poolside/Laguna-S-2.1-DFlash-NVFP4
            - name: HF_HOME
              value: /models/huggingface
            - name: HF_HUB_OFFLINE
              value: "0"
            - name: HF_HUB_DISABLE_XET
              value: "1"
            - name: HF_TOKEN
              valueFrom:
                secretKeyRef:
                  name: <HF_TOKEN_SECRET>
                  key: HF_TOKEN
                  optional: true
          volumeMounts:
            - name: model-cache
              mountPath: /models
      volumes:
        - name: model-cache
          persistentVolumeClaim:
            claimName: laguna-model-cache
---
# Headless Service — ordinal DNS for the Ray head address (MASTER_ADDR).
apiVersion: v1
kind: Service
metadata:
  name: laguna-tp2
  namespace: <NAMESPACE>
  labels:
    app: laguna-tp2
spec:
  clusterIP: None                 # headless — ordinal DNS for MASTER_ADDR (Ray head)
  publishNotReadyAddresses: true  # rank1 must resolve rank0 DNS before Ready
  selector:
    app: laguna-tp2
  ports:
    - name: http
      port: 8000
      targetPort: http
---
# rank0-only ClusterIP Service — the single API upstream (rank1 has no listener).
# Point your router / load balancer / InferencePool at THIS service.
apiVersion: v1
kind: Service
metadata:
  name: laguna-tp2-rank0
  namespace: <NAMESPACE>
  labels:
    app: laguna-tp2
spec:
  selector:
    app: laguna-tp2
    apps.kubernetes.io/pod-index: "0"
  ports:
    - name: http
      port: 8000
      targetPort: http
---
apiVersion: apps/v1
kind: StatefulSet
metadata:
  name: laguna-tp2
  namespace: <NAMESPACE>
  labels:
    app: laguna-tp2
spec:
  serviceName: laguna-tp2
  replicas: 2
  # Parallel REQUIRED: rank0's Ray head polls for rank1 to join the cluster
  # before launching vLLM; OrderedReady would deadlock (rank1 never created
  # while rank0 is not Ready). podManagementPolicy is immutable after creation.
  podManagementPolicy: Parallel
  selector:
    matchLabels:
      app: laguna-tp2
  template:
    metadata:
      labels:
        app: laguna-tp2
      annotations:
        # RoCE macvlan attached into the pod netns (as net1) for NCCL cross-node
        # all-reduce over the 200G ConnectX-7 fabric.
        k8s.v1.cni.cncf.io/networks: <ROCE_NAD>
    spec:
      nodeSelector:
        kubernetes.io/arch: arm64
      # Pin one pod per node across the two GB10 nodes (TP=2, 1 GPU/node).
      affinity:
        nodeAffinity:
          requiredDuringSchedulingIgnoredDuringExecution:
            nodeSelectorTerms:
              - matchExpressions:
                  - key: kubernetes.io/hostname
                    operator: In
                    values: [<NODE_A>, <NODE_B>]
        podAntiAffinity:
          requiredDuringSchedulingIgnoredDuringExecution:
            - labelSelector:
                matchLabels: {app: laguna-tp2}
              topologyKey: kubernetes.io/hostname
      terminationGracePeriodSeconds: 120
      securityContext:
        fsGroup: 0
      initContainers:
        - name: wait-model
          image: <REGISTRY>/vllm-spark-runtime:<IMAGE_TAG>
          command: [bash, -lc]
          args:
            - >-
              until test -f /models/.laguna_fetched
              && ls /models/huggingface/hub/models--poolside--Laguna-S-2.1-NVFP4/snapshots/*/config.json >/dev/null 2>&1
              && ls /models/huggingface/hub/models--poolside--Laguna-S-2.1-DFlash-NVFP4/snapshots/*/config.json >/dev/null 2>&1;
              do sleep 15; done
          env:
            - name: HF_HOME
              value: /models/huggingface
          volumeMounts:
            - name: model-cache
              mountPath: /models
      containers:
        - name: vllm
          # Ray backend (NOT mp — mp wedged this GB10 pair at 262144 AND 131072;
          # ray is the only no-skip-flag multi-node config that boots cleanly).
          # rank0 = Ray head + vLLM API leader; rank1 = `ray start --block`.
          # NOTE: the serve args below are a bash block, NOT a folded YAML scalar
          # with inline `#` comments (a whitespace-preceded `#` in a folded scalar
          # becomes a bash comment and silently truncates every following flag).
          image: <REGISTRY>/vllm-spark-runtime:<IMAGE_TAG>
          securityContext:
            runAsUser: 0
            capabilities:
              # RDMA memory registration pins pages; without IPC_LOCK
              # ibv_reg_mr_iova2 fails "Cannot allocate memory" -> NCCL error.
              add: [IPC_LOCK]
          command: [bash, -lc]
          args:
            - |
              if [ "${NODE_RANK}" = "0" ]; then
                ray start --head --port 6380 --num-cpus 8 \
                  --node-ip-address "${VLLM_HOST_IP}" \
                  --include-dashboard=false --disable-usage-stats &
                NODES=0
                for i in $(seq 1 60); do
                  NODES=$(ray status 2>/dev/null | grep -Ec '^\s*[0-9]+\s+node_' || true)
                  if [ "${NODES}" -ge 2 ]; then
                    echo "ray cluster at ${NODES} nodes; launching vllm serve"
                    break
                  fi
                  echo "waiting for rank1 to join ray cluster (attempt ${i}/60, nodes=${NODES})"
                  sleep 2
                done
                if [ "${NODES}" -lt 2 ]; then
                  echo "rank1 did not join ray cluster within ~120s; aborting" >&2
                  exit 1
                fi
                exec vllm serve poolside/Laguna-S-2.1-NVFP4 --served-model-name laguna-s-2.1 \
                  --host 0.0.0.0 --port 8000 \
                  --tensor-parallel-size 2 --pipeline-parallel-size 1 \
                  --distributed-executor-backend ray \
                  --speculative-config='{"model":"poolside/Laguna-S-2.1-DFlash-NVFP4","num_speculative_tokens":15,"method":"dflash"}' \
                  --max-model-len 1048576 --gpu-memory-utilization 0.80 --enforce-eager \
                  --kv-cache-memory 68719476736 \
                  --max-num-seqs 32 --max-num-batched-tokens 10240 --trust-remote-code \
                  --enable-prefix-caching --enable-prompt-tokens-details --enable-chunked-prefill \
                  --enable-auto-tool-choice --tool-call-parser poolside_v1 \
                  --reasoning-parser poolside_v1 --moe-backend flashinfer_b12x \
                  --default-chat-template-kwargs '{"enable_thinking":true}' \
                  --override-generation-config '{"temperature":0.7,"top_p":0.95}' \
                  --disable-access-log-for-endpoints /metrics
              else
                exec ray start --block --address="${MASTER_ADDR}:6380" \
                  --num-cpus 8 --disable-usage-stats \
                  --node-ip-address "${VLLM_HOST_IP}"
              fi
          env:
            - name: NODE_RANK
              valueFrom:
                fieldRef:
                  fieldPath: metadata.labels['apps.kubernetes.io/pod-index']
            - name: MASTER_ADDR
              value: laguna-tp2-0.laguna-tp2.<NAMESPACE>.svc.cluster.local
            - name: VLLM_HOST_IP
              valueFrom:
                fieldRef:
                  fieldPath: status.podIP
            - {name: HF_HOME, value: /models/huggingface}
            - {name: HF_HUB_OFFLINE, value: "1"}
            # flashinfer JIT does os.makedirs(~/.cache); point HOME at a writable dir.
            - {name: HOME, value: /opt/vllm}
            # b12x native NVFP4 MoE (FlashInferB12xExperts, CuTe DSL kernel path).
            # CUTE_DSL_ARCH selects the sm_121a CuTe DSL kernel variant at runtime.
            - {name: CUTE_DSL_ARCH, value: sm_121a}
            # Belt-and-suspenders in case a stale 256K config.json is still cached
            # (the fetcher normally patches config to the native 1M profile).
            - {name: VLLM_ALLOW_LONG_MAX_MODEL_LEN, value: "1"}
            - {name: MAX_JOBS, value: "4"}
            - {name: PATH, value: /usr/local/cuda/bin:/opt/env/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin}
            # expandable_segments reduces UMA fragmentation overshoot during the
            # memory-profiling forward pass.
            - {name: PYTORCH_CUDA_ALLOC_CONF, value: expandable_segments:True}
            # Ray executor knobs. memory_monitor_refresh_ms=0 stops Ray's own
            # memory monitor from competing with vLLM's budgeting during init.
            - {name: RAY_memory_monitor_refresh_ms, value: "0"}
            - {name: RAY_num_prestart_python_workers, value: "0"}
            - {name: RAY_object_store_memory, value: "1073741824"}
            # NCCL — pod netns: control plane on eth0; data plane via RoCE verbs.
            # GB10 has no GPUDirect RDMA; do NOT set NCCL_NET_GDR_LEVEL.
            # NCCL_IB_HCA must name the RoCE devices present in the pod netns.
            - {name: NCCL_NET, value: IB}
            - {name: NCCL_IB_DISABLE, value: "0"}
            - {name: NCCL_IB_HCA, value: "<ROCE_HCA_0>,<ROCE_HCA_1>"}
            - {name: NCCL_SOCKET_IFNAME, value: eth0}
            # Gloo reads its OWN iface var (not NCCL_SOCKET_IFNAME); without this
            # it can pick the host management NIC which is absent in the pod netns.
            - {name: GLOO_SOCKET_IFNAME, value: eth0}
            - {name: NCCL_CROSS_NIC, value: "1"}
            - {name: NCCL_CUMEM_ENABLE, value: "0"}
            - {name: NCCL_IGNORE_CPU_AFFINITY, value: "1"}
            - {name: NCCL_DEBUG, value: INFO}
            - {name: NCCL_NVLS_ENABLE, value: "0"}
            # NOTE: VLLM_SKIP_INIT_MEMORY_CHECK is intentionally NOT set — the
            # guard turns a UMA overshoot into a clean pod-fail rather than a
            # node wedge. This is WHY ray (not mp) is used here.
          ports:
            - {name: http, containerPort: 8000}
          # Probes branch on rank: rank1 runs `ray start --block` (no vLLM
          # process, no API listener), so rank1 health = "ray process alive".
          # Killing rank1 tears down the Ray cluster rank0 depends on.
          startupProbe:
            exec:
              command:
                - sh
                - -c
                - |
                  if [ "${NODE_RANK}" = "0" ]; then
                    python3 -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/v1/models', timeout=4)"
                  else
                    python3 -c "import glob,sys; sys.exit(0 if any(b'ray' in open(c,'rb').read() for c in glob.glob('/proc/[0-9]*/cmdline')) else 1)"
                  fi
            initialDelaySeconds: 30
            periodSeconds: 30
            timeoutSeconds: 10
            failureThreshold: 240   # generous headroom for the long 1M NVFP4 boot
          readinessProbe:
            exec:
              command:
                - sh
                - -c
                - |
                  if [ "${NODE_RANK}" = "0" ]; then
                    python3 -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/v1/models', timeout=3)"
                  else
                    python3 -c "import glob,sys; sys.exit(0 if any(b'ray' in open(c,'rb').read() for c in glob.glob('/proc/[0-9]*/cmdline')) else 1)"
                  fi
            periodSeconds: 10
            timeoutSeconds: 3
            failureThreshold: 3
          livenessProbe:
            exec:
              command:
                - sh
                - -c
                - |
                  if [ "${NODE_RANK}" = "0" ]; then
                    python3 -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health', timeout=8)"
                  else
                    python3 -c "import glob,sys; sys.exit(0 if any(b'ray' in open(c,'rb').read() for c in glob.glob('/proc/[0-9]*/cmdline')) else 1)"
                  fi
            periodSeconds: 30
            timeoutSeconds: 12
            failureThreshold: 10
          resources:
            # rdma/<...> is the RDMA device-plugin resource for the RoCE fabric.
            requests: {cpu: "8", memory: 105Gi, nvidia.com/gpu: "4", rdma/roce: "63"}
            limits:   {memory: 105Gi, nvidia.com/gpu: "4", rdma/roce: "63"}
          volumeMounts:
            - {name: shm, mountPath: /dev/shm}
            - {name: model-cache, mountPath: /models}
      volumes:
        - name: shm
          emptyDir:
            medium: Memory
            sizeLimit: 64Gi
        - name: model-cache
          persistentVolumeClaim:
            claimName: laguna-model-cache