[AMD][AgentX] kimik3: pin the KV page to 64 on an fp8+offload arm

fp8 KV + SimpleCPUOffloadConnector has now failed 3/3 (runs 30453589555,
30457683686; nodes g17/g19/g12), every time a GPU fault at under 10% pool
usage. The same connector on bf16 (run 30431115226) ran to 99.9% with zero
faults, and fp8 without the connector passes at c4/c8/c16. The two runs are
config-identical except kv_cache_dtype -- same vLLM g5f76ae224, GMU 0.88,
max-num-seqs 128, same cpu_bytes_to_use_per_rank -- and diffing the two
vllm_command.txt files leaves --kv-transfer-config as the only delta.

Both allocate an identical 14122 CPU blocks while GPU token capacity goes
1,957,290 -> 3,859,490. manager.py:199 scales num_cpu_blocks with
num_gpu_blocks, so the block count held and tokens/block doubled: fp8 takes a
128-token page where bf16 takes 64. A 128 page routes MLA into aiter's Gluon
b128 kernel, which fp8 survives on its own but has not survived alongside
block-granular connector copies.

Add KV_BLOCK_SIZE (-> --block-size) and a vllm-simple-fp8-b64 backend that
pins it to 64, plus a one-cell fp8b64 key. Keeps fp8 and moves exactly one
variable off the failing config.

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

Author: hyukjlee

benchmarks/single_node/agentic/kimik3_fp4_mi355x.sh

@@ -46,6 +46,7 @@ set -x
 #   AITER_A8W4               1      (reference; 0 = aiter a16w4 MoE path)
 #   LANGUAGE_MODEL_ONLY      false  (reference loads the vision tower)
 #   KV_CACHE_DTYPE           fp8    (default for every arm; =auto for a bf16 A/B)
+#   KV_BLOCK_SIZE            unset  (unset -> vLLM sizes the page; 128 under fp8)
 #   MAX_MODEL_LEN            unset  (unset -> vLLM derives K3's 1M context)
 #   SPEC_DECODE              false  (enabled by the _mtp DSpark wrapper)
 #   ENFORCE_EAGER            false  (enabled by the _mtp DSpark wrapper)
@@ -427,18 +428,30 @@ except Exception:
         )
     fi
     ;;
-  vllm-simple|vllm-simple-fp8)
+  vllm-simple|vllm-simple-fp8|vllm-simple-fp8-b64)
     require_agentic_kv_offload_backend "$KV_OFFLOAD_BACKEND"
+    # vllm-simple-fp8-b64 is vllm-simple-fp8 with the KV page pinned to 64
+    # tokens. Under fp8 vLLM otherwise picks 128, which routes MLA into aiter's
+    # Gluon b128 kernel; fp8 survives that alone but every fp8 cell that also
+    # ran this connector died with a GPU fault (3/3, g17/g19/g12, at <10% pool
+    # usage) while bf16 + the same connector ran to 99.9%. 64 is the page bf16
+    # used, so this holds the one geometry that is known to work with the
+    # connector while keeping fp8.
+    if [ "$KV_OFFLOAD_BACKEND" = "vllm-simple-fp8-b64" ]; then
+        KV_BLOCK_SIZE="${KV_BLOCK_SIZE:-64}"
+    fi
     # vllm-simple-fp8 is the same connector with an fp8 KV cache. fp8 halves
     # bytes/token in the GPU pool, which on this KV-bound corpus moves the
     # eviction wall itself rather than just adding headroom (measured: the pool
     # peaks at 98.6-99.8% usage even at c8). ROCm maps fp8 -> fp8_e4m3.
     # UNVALIDATED on K3's hybrid geometry: the pool spans Kimi Delta Attention
     # state and gated-MLA latent, and fp8 support across both spec types is
     # unconfirmed on this build.
-    if [ "$KV_OFFLOAD_BACKEND" = "vllm-simple-fp8" ]; then
+    case "$KV_OFFLOAD_BACKEND" in
+      vllm-simple-fp8|vllm-simple-fp8-b64)
         KV_CACHE_DTYPE="${KV_CACHE_DTYPE:-fp8}"
-    fi
+        ;;
+    esac
     # vLLM's own SimpleCPUOffloadConnector -- the AMD reference's native
     # offload path. Unlike LMCache it does not go through the Mamba
     # [conv_state, ssm_state] adapter that K3's Kimi Delta Attention breaks
@@ -537,6 +550,22 @@ if [ -n "${KV_CACHE_DTYPE:-}" ] && [ "${KV_CACHE_DTYPE}" != "auto" ]; then
     KV_CACHE_DTYPE_ARGS=(--kv-cache-dtype "$KV_CACHE_DTYPE")
 fi
 
+# Unset by default: vLLM sizes the KV page itself. Under fp8 it picks a 128-token
+# page (bf16 gets 64) -- measured as an identical 14122 CPU blocks in the offload
+# connector across both dtypes while GPU token capacity went 1,957,290 ->
+# 3,859,490, which per manager.py:199 (num_cpu_blocks scales with num_gpu_blocks)
+# can only mean the block count held and tokens/block doubled.
+#
+# A 128-token page routes MLA into aiter's Gluon b128 kernel. That path is fine
+# on its own (the gist patch above; fp8 kvnone cells pass), but every fp8 cell
+# that ALSO ran SimpleCPUOffloadConnector has died with a GPU fault -- 3/3 on
+# g17/g19/g12, at <10% pool usage, while the same connector on bf16 ran to 99.9%.
+# Set KV_BLOCK_SIZE=64 to hold fp8 on the geometry bf16 used.
+KV_BLOCK_SIZE_ARGS=()
+if [ -n "${KV_BLOCK_SIZE:-}" ] && [ "${KV_BLOCK_SIZE}" != "0" ]; then
+    KV_BLOCK_SIZE_ARGS=(--block-size "$KV_BLOCK_SIZE")
+fi
+
 # Left unset by default so vLLM derives K3's native 1M context, which is what
 # the unfiltered corpus needs. Set explicitly only to test truncation effects.
 MAX_MODEL_LEN_ARGS=()
@@ -719,6 +748,7 @@ VLLM_CMD=(
     "${MAX_MODEL_LEN_ARGS[@]}"
     "${PREFIX_CACHE_ARGS[@]}"
     "${KV_CACHE_DTYPE_ARGS[@]}"
+    "${KV_BLOCK_SIZE_ARGS[@]}"
     "${EAGER_ARGS[@]}"
     "${SPEC_ARGS[@]}"
     "${EVAL_SERVE_ARGS[@]}"

configs/amd-master.yaml

@@ -775,6 +775,41 @@ kimik3-fp4-mi355x-vllm-agentic-fp8retry:
       - { tp: 8, kv-offloading: none, conc-list: [1, 16] }
       - { tp: 8, kv-offloading: dram, kv-offload-backend: { name: vllm-simple-fp8 }, conc-list: [8] }
 
+kimik3-fp4-mi355x-vllm-agentic-fp8b64:
+  # One variable off the cell that has now failed 3/3: fp8 KV + the
+  # SimpleCPUOffloadConnector, with the KV page pinned to 64 tokens.
+  #
+  # Runs 30453589555 and 30457683686 killed every fp8+connector cell with a GPU
+  # fault (g17/g19/g12) at under 10% pool usage, while the SAME connector on
+  # bf16 (run 30431115226) ran to 99.9% with zero faults, and fp8 WITHOUT the
+  # connector passes (kvnone c4/c8/c16). The two runs are config-identical
+  # except kv_cache_dtype -- same vLLM g5f76ae224, GMU 0.88, max-num-seqs 128,
+  # same cpu_bytes_to_use_per_rank -- and diffing vllm_command.txt leaves
+  # --kv-transfer-config as the only delta.
+  #
+  # Both runs allocate an identical 14122 CPU blocks while GPU token capacity
+  # goes 1,957,290 -> 3,859,490. manager.py:199 scales num_cpu_blocks with
+  # num_gpu_blocks, so the block count held and tokens/block doubled: fp8 picks
+  # a 128-token page, bf16 gets 64. A 128 page routes MLA into aiter's Gluon
+  # b128 kernel -- survivable alone (the gist patch; run 30442578333 got 696
+  # tok/s/GPU) but not, so far, alongside block-granular connector copies.
+  #
+  # If this passes, the b128 page is the culprit and fp8 keeps the host tier.
+  # If it dies the same way, the page is not the variable and the next step is
+  # reading the worker-side copy path rather than guessing at geometry.
+  image: vllm/vllm-openai-rocm:kimi-k3
+  model: moonshotai/Kimi-K3
+  model-prefix: kimik3
+  runner: cluster:mi355x-amds
+  precision: fp4
+  framework: vllm
+  multinode: false
+  scenarios:
+    agentic-coding:
+    - dram-utilization: 0.80
+      search-space:
+      - { tp: 8, kv-offloading: dram, kv-offload-backend: { name: vllm-simple-fp8-b64 }, conc-list: [8] }
+
 kimik3-fp4-mi355x-vllm-agentic-fp8kv:
   # fp8 KV on the OFFLOAD arm. This is the fp8kvtest arm (run 30442578333:
   # 696 total tok/s/GPU, TTFT 2.6s, the first clean fp8 result) plus the