Skip to content

Instantly share code, notes, and snippets.

@co-l
Created August 23, 2026 12:04
Show Gist options
  • Select an option

  • Save co-l/c2aeaf40b53fcacfe9dd3293be75f23a to your computer and use it in GitHub Desktop.

Select an option

Save co-l/c2aeaf40b53fcacfe9dd3293be75f23a to your computer and use it in GitHub Desktop.
Qwen3.8-27B NVFP4 with vision + 451K token KV-cache on one RTX 5090 (400W) — standalone reproduction

Qwen 3.8 27B NVFP4 with vision + 451K tokens KV-cache on one RTX 5090@400W: 120 tokens/s

What this is

Qwen3.8-27B, with vision, holding 451K tokens of KV on a single RTX 5090, under a 400 W cap, in a container. Stock fp8 KV fits ~124K total on this card; this stack does 3.6x that, so three parallel sessions' context (3×75K) stays resident in the pool: sessions keep their KV and continue without re-prefill cost.

The stack

Component Choice Why
Model gittensor-model-hub/Qwen3.8-27B-NVFP4-RTX5090 Model Optimizer NVFP4, everything quantized to FP4 → smallest weights (~15 GB) → biggest KV budget
Engine patched vLLM nightly (localhost/vllm-nvfp4kv:latest) stock vLLM gates NVFP4 KV to SM100; patch routes it through FlashInfer's FA2 reader on SM120
KV cache nvfp4 + HND layout 4-bit KV ≈ 3.5x fp8 capacity; HND layout is non-negotiable
Spec decode MTP, num_speculative_tokens=3, in-checkpoint the model's own MTP heads, weights shared, zero extra VRAM
Cudagraphs PIECEWISE mandatory; see the ONE gotcha below
Context 196,608 the model's native length is 262,144; 196,608 is chosen to maximize concurrency: smaller average per-session context means more parallel sessions fit in the 451K pool
Power 400 W + 300–2500 MHz + VF offset +600 (LACT) quiet, efficient
GPU util / seqs 0.98 / 4 max KV pool, 4 concurrent requests

Why this model (and not the others)

All the NVFP4 Qwen3.8-27B checkpoints on HF reduce to three recipes. This one picks the aggressive end on purpose: the smaller the weights, the more KV fits:

Recipe Weights What stays high-precision KV pool @196K Verdict
gittensor (ModelOpt) 15 GB nothing (FP4 everywhere incl. attention) 451K this stack: smallest, fastest, top-scoring
unsloth (compressed-tensors) 20 GB FP8 attention 270K fine, but 1.7x less context for no quality edge
RadixArk (ModelOpt mixed) 21 GB FP8 attention n/a SGLang-first; does not load in vLLM
AutoRound family ~16 GB dequant kernels 395K good quality, ~3x slower prefill
GDN-kept-BF16 family (sakamakismile, NVFP4A16) 27 GB GDN state in BF16 ~180K best precision, but eats the KV budget, can't hold 3×75K

Benchmarks

400 W, MTP-3, PIECEWISE, NVFP4 KV. Short-context rows (4K/8K) measured 2026-08-23 on this exact build; long-context rows from the 2026-08-22 validation (identical stack). pp = prefill (t/s), tg64 = 64-token decode (t/s), ttfr = time to first token (ms).

Single session

prompt pp t/s ttfr (ms) tg64 t/s
4K 11,388 352 130
8K 10,475 765 150
16K 8,750 1,830 126
32K 6,790 4,714 131
64K 4,729 13,535 111
128K 2,932 43,661 107
185K 2,194 84,306 112

Decode holds ~107–150 t/s at any context. Prefill decays with context (attention becomes memory-bound over the whole KV), which is normal.

Concurrency (the real workload)

test agg t/s per-req t/s ttfr (ms)
pp4K c2 11,856 6,653 610
pp4K c3 11,795 4,907 848
pp8K c2 10,648 6,240 1,311
pp8K c3 10,637 4,779 1,786
tg64 c2 (4K) 199 118 n/a
tg64 c3 (4K) 204 99 n/a
tg64 c3 (8K) 124 86 n/a

Prefill scales down per-request as expected; decode stays usable per-session under load. Minor prefill hiccups are sometimes visible when several sessions prefill at once, nothing harmful.

The target workload: 3 sessions × 75K

  • Total context: 225K of a 451K pool → 50%. Fits comfortably, with room for a fourth session at short context.
  • Decode at concurrency 3: ~86–99 t/s per session.
  • Prefill at concurrency 3: ~4.9K t/s per request.

Real coding session (validated 2026-08-22, same stack)

An agent wrote a small game while spawning 2 sub-agents, with vision checks mixed in. 6 responses, context growing to 116K tokens:

# ctx (k, start → end) calls TG t/s
1 8.7 → 14.2 10 136.4
2 15.4 → 101.8 42 135.5
3 3.8 → 28.0 7 144.6
4 3.0 → 28.1 8 120.8
5 103.1 → 115.5 14 112.6
6 → 116.0 1 109.9

Decode held 109–145 t/s through a 116K context, matching the synthetic benches. This is the config's proof of life for real agentic work.

The ONE gotcha (read this before touching anything)

Never use FULL cudagraphs with MTP spec-decode + NVFP4 KV. Use PIECEWISE.

FULL cudagraph capture succeeds: the log even says "Capturing CUDA graphs (decode, FULL): 100%", but replay silently corrupts the model's reasoning stream: outputs degenerate into loops, empty answers, or hallucinated content from other requests. It does not crash, so it is easy to mistake for a bad model (that is exactly what happened to us; an innocent model nearly took the blame).

The vLLM launch path auto-selects PIECEWISE when NVFP4 KV + spec-decode are active; do not override it. PIECEWISE costs a little per-step launch overhead versus FULL (~1-2% on this dense model), a price worth paying for correct output.


PART 2: Reproduction (for agents)

This half is written to be reproduced, not explained.

0. Environment (exactly what we run)

Component Version
GPU RTX 5090 32 GiB, compute 12.0 (SM120), 400 W cap
Host Bazzite (Fedora-based), 30 GiB RAM, rootful podman, headless
Base image docker.io/vllm/vllm-openai:nightly = 0.26.1rc1.dev1102+ge9d1398d9 (2026-08-22)
flashinfer 0.6.17 (bundled with the nightly)
torch / nvcc 2.13.0+cu130 / 13.0 (inside the base image)
Model gittensor-model-hub/Qwen3.8-27B-NVFP4-RTX5090 (multimodal)
Power LACT: 400 W cap, 300–2500 MHz, VF offset +600

The compile step runs inside the base image (it ships nvcc + torch headers); no CUDA toolkit on the host.

1. Why the image is patched

Stock vLLM refuses --kv-cache-dtype nvfp4 on the 5090 because nvfp4 KV is gated to SM100+ (the trtllm-gen FP4 kernels only exist there). The SM120 kernels already exist in flashinfer (the FA2 paged NVFP4 reader), vLLM just never routes to them. Two fixes, no vLLM source rebuild:

  1. The SM100 gate: flashinfer.py accepts nvfp4 on SM120 and routes it through the FA2 paged reader (model-dtype queries/output, head-major HND layout).
  2. The V-scale swizzle bug: the stock writer stores V block scales in the SM100 trtllm-gen 4-token swizzle unconditionally; the FA2 reader expects them linear. Result if unfixed: silent corruption (!!!!!!-flooded output). A standalone ~300-line kernel (my_nvfp4_kv.so) writes linear V scales on SM12x.

All sources needed are embedded verbatim in the self-contained builder below (build-nvfp4kv-image.sh): the flashinfer.py unified diff, the standalone writer kernel, its two vendored headers, and the compile driver. Save that script and run it; nothing else is required.

2. Build the image (self-contained)

Save the script below as build-nvfp4kv-image.sh, make it executable, and run it:

chmod +x build-nvfp4kv-image.sh && ./build-nvfp4kv-image.sh

What it does: materializes the five embedded sources, ships them to the box, compiles my_nvfp4_kv.so (sm_120f) in a throwaway container, applies the flashinfer.py diff (installing patch in the throwaway container if needed), and commits localhost/vllm-nvfp4kv:latest from the nightly with the patched file + the kernel .so. Idempotent.

Compile landmines (already handled in the embedded build_ext.py):

  • CUDA headers live in the pip nvidia/cu13 package -> CUDA_HOME points there
  • CCCL toolkit compatibility check -> -DCCCL_DISABLE_CTK_COMPATIBILITY_CHECK
  • __half->float conversion (CUDA 13 has no implicit one) -> explicit helper
  • fp4 needs sm_120f (family-conditional) -> monkeypatched arch flags
  • libcudart.so has no unversioned symlink -> created before linking
  • TORCH_LIBRARY alone produces no PyInit_ -> PYBIND11_MODULE added

podman commit bakes the throwaway --entrypoint /bin/bash unless overridden; the script sets ENTRYPOINT ["vllm", "serve"] explicitly.

#!/bin/bash
# Self-contained builder for the SM120 NVFP4-KV vLLM image.
# Creates localhost/vllm-nvfp4kv:latest on $BOX. Every source is embedded
# below; nothing else is required. Run from any machine that can ssh to the
# box and has rootful podman. Idempotent.
set -euo pipefail

BOX="${BOX:-gaming_pc}"
SRC=/tmp/nvfp4standalone
BOX_SRC=/tmp/nvfp4build

# 1. materialize the embedded sources locally
rm -rf "$SRC" && mkdir -p "$SRC"
cat > "$SRC/flashinfer.py.diff" <<'PATCH_EOF'
--- /tmp/nvfp4-patch/flashinfer_nightly.py	2026-08-22 14:04:38.044735256 +0300
+++ /tmp/nvfp4-patch/flashinfer.py	2026-08-22 13:28:37.992822504 +0300
@@ -38,6 +38,11 @@
 )
 from vllm.platforms import current_platform
 from vllm.platforms.interface import DeviceCapability
+
+try:
+    import my_nvfp4_kv  # noqa: F401  # registers torch.ops._my_nvfp4_kv
+except ImportError:
+    my_nvfp4_kv = None
 from vllm.triton_utils import tl, triton
 from vllm.utils.flashinfer import (
     can_use_trtllm_attention,
@@ -480,6 +485,11 @@
     @classmethod
     def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool:
         if kv_cache_dtype is not None and kv_cache_dtype.startswith("nvfp4"):
+            if current_platform.is_device_capability_family(120):
+                # Consumer Blackwell (sm120/sm121): NVFP4 KV is served through
+                # the FlashInfer FA2 paged reader (uint8 fp4 cache), no
+                # trtllm-gen requirement.
+                return True
             return (
                 current_platform.is_device_capability_family(100)
                 and supports_trtllm_attention(is_prefill=True)
@@ -530,6 +540,18 @@
             # The trtllm-gen kernels consume head-major block interiors; the L/B
             # nesting outside the block is immaterial to them.
             return (KVCacheLayout.LBHNC, KVCacheLayout.BLHNC)
+        if capability is not None and capability.major == 12:
+            # NVFP4 KV on consumer Blackwell (sm120/sm121, FA2 path): each K/V
+            # side packs [data | scale] regions carved out of the side's byte
+            # range, which is only byte-coherent under the head-major HND
+            # layout (LBHNC). NHD would silently corrupt the cache.
+            vllm_config = get_current_vllm_config_or_none()
+            if (
+                vllm_config is not None
+                and vllm_config.cache_config is not None
+                and (vllm_config.cache_config.cache_dtype or "").startswith("nvfp4")
+            ):
+                return (KVCacheLayout.LBHNC,)
         return super().supported_kv_cache_layouts()
 
     forward_includes_kv_cache_update: bool = False
@@ -753,8 +775,19 @@
             # Cannot use self.kv_cache_spec.dtype here because kv_cache_spec
             # storage dtype may not be the same as the op dtype (uint8 vs fp8_e4m3)
             self.is_kvcache_nvfp4 = self.cache_dtype.startswith("nvfp4")
+            self.use_fa2_nvfp4_kv = False
             if self.is_kvcache_nvfp4:
-                if (
+                if current_platform.is_device_capability_family(120):
+                    # Consumer Blackwell (sm120/sm121): no trtllm-gen FP4 FMHA,
+                    # so route NVFP4 KV through FlashInfer's FA2 paged reader.
+                    # The cache stores packed uint8 fp4 data; the per-side
+                    # [data | scale] regions are read back as views with
+                    # explicit strides (nvfp4_split_data_scale).
+                    self.use_fa2_nvfp4_kv = True
+                    self.kv_cache_dtype = FlashInferBackend.get_dtype_for_flashinfer(
+                        "nvfp4"
+                    )
+                elif (
                     force_use_trtllm_attention() is False
                     or not supports_trtllm_attention(is_prefill=True)
                     or not supports_trtllm_attention(is_prefill=False)
@@ -764,9 +797,10 @@
                         "SM100 trtllm-gen "
                         "FlashInfer path."
                     )
-                # The scale search only affects the store kernel. FlashInfer
-                # reads both variants using the same NVFP4 layout.
-                self.kv_cache_dtype = "nvfp4"
+                else:
+                    # The scale search only affects the store kernel. FlashInfer
+                    # reads both variants using the same NVFP4 layout.
+                    self.kv_cache_dtype = "nvfp4"
             else:
                 self.kv_cache_dtype = FlashInferBackend.get_dtype_for_flashinfer(
                     self.cache_dtype
@@ -774,6 +808,7 @@
         else:
             self.cache_dtype = "auto"
             self.is_kvcache_nvfp4 = False
+            self.use_fa2_nvfp4_kv = False
             assert self.kv_cache_spec.dtype == self.model_config.dtype
             self.kv_cache_dtype = self.kv_cache_spec.dtype
 
@@ -802,6 +837,12 @@
             if can_use_xqa_or_trtllm_gen_decode
             else None
         )
+        if self.use_fa2_nvfp4_kv:
+            # NVFP4 KV on consumer Blackwell uses the FlashInfer FA2 paged
+            # reader for decode (flashinfer-native); XQA/trtllm-gen cannot
+            # consume the nvfp4 cache, so keep decode on the FA2 wrapper.
+            self.use_trtllm_decode_attention = False
+            self.flashinfer_trtllm_api_decode_kernel = None
         # The dedicated FlashInfer XQA API accepts head dimensions in
         # [16, 256] that are divisible by 16. Some hybrid-attention models
         # (for example Gemma 4) use XQA-compatible sliding-attention groups
@@ -947,6 +988,10 @@
                 return FlashInferBackend.get_dtype_for_flashinfer(cache_dtype)
             return self.model_config.dtype
         if cache_dtype.startswith("nvfp4"):
+            if self.use_fa2_nvfp4_kv:
+                # The FA2 paged nvfp4 reader (consumer Blackwell) consumes
+                # model-dtype queries; FP8-Q is a trtllm-gen-only contract.
+                return self.model_config.dtype
             return FlashInferBackend.get_dtype_for_flashinfer("fp8_e4m3")
         return self.kv_cache_spec.dtype
 
@@ -1107,13 +1152,14 @@
                 raise NotImplementedError(
                     "FlashInfer non-causal prefill is not supported with DCP yet."
                 )
-            if self.is_kvcache_nvfp4:
+            if self.is_kvcache_nvfp4 and not self.use_fa2_nvfp4_kv:
                 raise NotImplementedError(
                     "FlashInfer non-causal attention is not supported with "
                     "NVFP4 KV cache."
                 )
             if self._noncausal_prefill_wrapper is None:
                 if self.has_sinks and current_platform.is_device_capability_family(120):
+                    assert not self.is_kvcache_nvfp4
                     self._noncausal_prefill_wrapper = (
                         BatchAttentionWithAttentionSinkWrapper(
                             self._get_workspace_buffer(),
@@ -1127,11 +1173,14 @@
                         )
                     )
                 else:
+                    # NVFP4 KV on the FA2 path (sm12x) supports non-causal
+                    # prefill via the fa2 backend; trtllm-gen/auto cannot.
+                    backend = "fa2" if self.use_fa2_nvfp4_kv else "auto"
                     self._noncausal_prefill_wrapper = (
                         BatchPrefillWithPagedKVCacheWrapper(
                             self._get_workspace_buffer(),
                             get_flashinfer_layout_string(self.kv_cache_layout),
-                            backend="auto",
+                            backend=backend,
                         )
                     )
             return self._noncausal_prefill_wrapper
@@ -1157,9 +1206,12 @@
                         window_left=self.window_left,
                     )
                 else:
-                    # NVFP4 KV cache requires the trtllm-gen backend inside
-                    # the wrapper; fa2/fa3 do not support nvfp4.
-                    backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto"
+                    # NVFP4 KV: FlashInfer FA2 paged reader on consumer
+                    # Blackwell (sm120/sm121); trtllm-gen on sm100f.
+                    if self.use_fa2_nvfp4_kv:
+                        backend = "fa2"
+                    else:
+                        backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto"
                     self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(
                         self._get_workspace_buffer(),
                         get_flashinfer_layout_string(self.kv_cache_layout),
@@ -1183,9 +1235,12 @@
                 paged_kv_indptr = None
                 paged_kv_indices = None
                 paged_kv_last_page_len = None
-            # NVFP4 KV cache requires the trtllm-gen backend inside
-            # the wrapper; fa2/fa3 do not support nvfp4.
-            backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto"
+            # NVFP4 KV: FlashInfer FA2 paged reader on consumer Blackwell
+            # (sm120/sm121); trtllm-gen on sm100f.
+            if self.use_fa2_nvfp4_kv:
+                backend = "fa2"
+            else:
+                backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto"
             decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
                 self._get_workspace_buffer(),
                 get_flashinfer_layout_string(self.kv_cache_layout),
@@ -1592,7 +1647,9 @@
                     # use FP8 o_data_type so the wrapper matches the
                     # FP8 output buffer allocated in forward().
                     o_dtype = (
-                        FP8_DTYPE if self.is_kvcache_nvfp4 else self.model_config.dtype
+                        FP8_DTYPE
+                        if (self.is_kvcache_nvfp4 and not self.use_fa2_nvfp4_kv)
+                        else self.model_config.dtype
                     )
                     prefill_wrapper.plan(
                         qo_indptr=qo_indptr_prefill_cpu,
@@ -1677,7 +1734,9 @@
                 # use FP8 o_data_type so the wrapper matches the
                 # FP8 output buffer allocated in forward().
                 o_dtype = (
-                    FP8_DTYPE if self.is_kvcache_nvfp4 else self.model_config.dtype
+                    FP8_DTYPE
+                    if (self.is_kvcache_nvfp4 and not self.use_fa2_nvfp4_kv)
+                    else self.model_config.dtype
                 )
                 fast_plan_decode(
                     decode_wrapper,
@@ -1747,6 +1806,9 @@
         )
         self.cache_dtype = kv_cache_dtype
         self.is_kvcache_nvfp4 = kv_cache_dtype.startswith("nvfp4")
+        self.use_fa2_nvfp4_kv = (
+            self.is_kvcache_nvfp4 and current_platform.is_device_capability_family(120)
+        )
         self.kv_cache_dtype = "nvfp4" if self.is_kvcache_nvfp4 else kv_cache_dtype
         self.fp4_data_dim = head_size // 2 if self.is_kvcache_nvfp4 else 0
         self.logits_soft_cap = logits_soft_cap
@@ -1795,7 +1857,12 @@
         self.o_sf_scale: float | None = None
 
         # Pre-allocated FP8 output buffer for NVFP4 without fused output quant.
-        if self.is_kvcache_nvfp4 and vllm_config is not None:
+        # The sm12x FA2 path writes into `output` directly and never reads this.
+        if (
+            self.is_kvcache_nvfp4
+            and vllm_config is not None
+            and not self.use_fa2_nvfp4_kv
+        ):
             max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
             self._nvfp4_fp8_out = torch.empty(
                 (max_num_tokens, num_heads, head_size),
@@ -2118,7 +2185,9 @@
                     # Use a pre-allocated FP8 buffer and dequantize
                     # afterwards.
                     needs_fp8_out_prefill = (
-                        self.is_kvcache_nvfp4 and output.dtype != FP8_DTYPE
+                        self.is_kvcache_nvfp4
+                        and not self.use_fa2_nvfp4_kv
+                        and output.dtype != FP8_DTYPE
                     )
                     if needs_fp8_out_prefill:
                         out_prefill = self._nvfp4_fp8_out[:num_prefill_tokens]
@@ -2185,7 +2254,11 @@
 
                 # NVFP4 trtllm kernel only supports FP8 output.
                 # Use a pre-allocated FP8 buffer and dequantize afterwards.
-                needs_fp8_out = self.is_kvcache_nvfp4 and output.dtype != FP8_DTYPE
+                needs_fp8_out = (
+                    self.is_kvcache_nvfp4
+                    and not self.use_fa2_nvfp4_kv
+                    and output.dtype != FP8_DTYPE
+                )
                 if needs_fp8_out:
                     out = self._nvfp4_fp8_out[:num_prefill_tokens]
 
@@ -2288,7 +2361,11 @@
 
                 # NVFP4 kernel only supports FP8 output.
                 # Use a pre-allocated FP8 buffer and dequantize afterwards.
-                needs_fp8_out = self.is_kvcache_nvfp4 and output.dtype != FP8_DTYPE
+                needs_fp8_out = (
+                    self.is_kvcache_nvfp4
+                    and not self.use_fa2_nvfp4_kv
+                    and output.dtype != FP8_DTYPE
+                )
                 if needs_fp8_out:
                     out_decode = self._nvfp4_fp8_out[:num_decode_tokens]
                 else:
@@ -2417,7 +2494,11 @@
 
                 # NVFP4 trtllm kernel only supports FP8 output.
                 # Use a pre-allocated FP8 buffer and dequantize afterwards.
-                needs_fp8_out = self.is_kvcache_nvfp4 and output.dtype != FP8_DTYPE
+                needs_fp8_out = (
+                    self.is_kvcache_nvfp4
+                    and not self.use_fa2_nvfp4_kv
+                    and output.dtype != FP8_DTYPE
+                )
                 if needs_fp8_out:
                     out = self._nvfp4_fp8_out[:num_decode_tokens]
 
@@ -2519,6 +2600,26 @@
                 k_cache, v_cache = kv_cache.transpose(1, 2).split(
                     self.head_size, dim=-1
                 )
+            if self.use_fa2_nvfp4_kv:
+                # Consumer Blackwell FA2 path: write V scales linearly (the
+                # stock _C_stable_libtorch writer always uses the SM100
+                # trtllm-gen swizzle, which corrupts the FA2 read). Use the
+                # standalone fixed writer.
+                if my_nvfp4_kv is None:
+                    raise RuntimeError(
+                        "NVFP4 KV on SM120 requires the my_nvfp4_kv extension "
+                        "(linear V scale-factor writer); it is not installed."
+                    )
+                my_nvfp4_kv.reshape_and_cache_nvfp4(
+                    key,
+                    value,
+                    k_cache,
+                    v_cache,
+                    slot_mapping,
+                    layer._k_scale,
+                    layer._v_scale,
+                )
+                return
             torch.ops._C_cache_ops.reshape_and_cache_flash(
                 key,
                 value,
PATCH_EOF

cat > "$SRC/my_nvfp4_kv.cu" <<'CU_EOF'
// Standalone NVFP4 KV cache writer for consumer Blackwell (sm120/sm121).
//
// Backport of the upstream vLLM nvfp4 KV store kernel (csrc/libtorch_stable/
// nvfp4_kv_cache_kernels.cu) with the V-scale swizzle fix: on SM12x the V
// block scales are written linearly (the FlashInfer FA2 paged reader reads
// scale factors from the SF tensor with linear strides), while SM100 keeps
// the trtllm-gen 4-token swizzle.
//
// Compiled as a tiny standalone torch extension so the full
// `_C_stable_libtorch` extension does not need rebuilding.

#include <torch/extension.h>
#include <torch/csrc/utils/pybind.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_fp8.h>
#include <cuda_runtime.h>

#include <algorithm>
#include <cstdint>
#include <string>
#include <utility>

#include "nvfp4_utils.cuh"
#include "cuda_vec_utils.cuh"

namespace vllm {

// Compute swizzled scale offset for SM100 trtllm-gen MHA kernel.
// The swizzle pattern for HND layout is:
//   [T//4, 4, 4, S//4] -> permute(0, 2, 3, 1) -> reshape to [T, S]
// where T = block_size (page_size), S = scale_dim = head_size // 16.
//
// For a linear (t, s) position, the swizzled position is:
//   swizzled_t = (t / 4) * 4 + (s / (S / 4))
//   swizzled_s = (s % (S / 4)) * 4 + (t % 4)
__device__ __forceinline__ int swizzle_scale_offset(int t, int s,
                                                    int scale_dim) {
  int s_group = scale_dim / 4;
  int swizzled_t = (t / 4) * 4 + (s / s_group);
  int swizzled_s = (s % s_group) * 4 + (t % 4);
  return swizzled_t * scale_dim + swizzled_s;
}

// Kernel: quantize bf16/fp16 key/value to NVFP4 and store in paged KV cache.
//
// Per page layout: [K_data | K_scale | V_data | V_scale]. Within each KV
// side, data and scale are separate contiguous regions.
//
// Threading: one CUDA block per token, threads process heads and
// groups of 16 elements within each head.
template <typename scalar_t>
__global__ void reshape_and_cache_nvfp4_kernel(
    const scalar_t* __restrict__ key,      // [num_tokens, num_heads, head_size]
    const scalar_t* __restrict__ value,    // [num_tokens, num_heads, head_size]
    uint8_t* __restrict__ key_data_cache,  // data region for K
    uint8_t* __restrict__ value_data_cache,    // data region for V
    uint8_t* __restrict__ key_scale_cache,     // scale region for K
    uint8_t* __restrict__ value_scale_cache,   // scale region for V
    const int64_t* __restrict__ slot_mapping,  // [num_actual_tokens]
    const float* __restrict__ k_scale_ptr,     // pointer to checkpoint k_scale
    const float* __restrict__ v_scale_ptr,     // pointer to checkpoint v_scale
    const int64_t key_stride,                  // key.stride(0) in elements
    const int64_t value_stride,                // value.stride(0) in elements
    const int num_heads, const int head_size, const int block_size,
    const int64_t data_block_stride,         // data cache stride for dim 0
    const int64_t data_head_stride,          // data cache stride for heads
    const int64_t data_block_offset_stride,  // data cache stride for tokens
    const int64_t scale_block_stride,        // scale cache stride for dim 0
    const int64_t scale_head_stride,         // scale cache stride for heads
    const int64_t scale_block_offset_stride,  // scale cache stride for tokens
    const bool swizzle_v_sf  // V scale layout: true = SM100 trtllm-gen 4-token
                             // swizzle; false = linear (FlashInfer FA2 sm12x)
) {
  using CudaType = typename CUDATypeConverter<scalar_t>::Type;
  using PVec = PackedVec<CudaType, CVT_FP4_PACK16>;

  static constexpr int ELTS = CVT_FP4_ELTS_PER_THREAD;  // 16 or 8
  static constexpr int THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / ELTS;

  const int64_t token_idx = blockIdx.x;
  const int64_t slot_idx = slot_mapping[token_idx];
  if (slot_idx < 0) return;

  const int64_t block_idx = slot_idx / block_size;
  const int block_offset = static_cast<int>(slot_idx % block_size);

  const int scale_dim = head_size / 16;
  const int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE;

  const int total_groups = num_heads * groups_per_head;
  const int tid = threadIdx.x;
  const int num_thread_groups = blockDim.x / THREADS_PER_SF;
  const int tg_id = tid / THREADS_PER_SF;
  const int tg_lane = tid % THREADS_PER_SF;

  // Process both K (kv=0) and V (kv=1)
#pragma unroll
  for (int kv = 0; kv < 2; kv++) {
    const scalar_t* __restrict__ src = (kv == 0) ? key : value;
    const float global_scale = 1.0f / ((kv == 0) ? *k_scale_ptr : *v_scale_ptr);
    const int64_t src_stride = (kv == 0) ? key_stride : value_stride;
    uint8_t* __restrict__ data_cache =
        (kv == 0) ? key_data_cache : value_data_cache;
    uint8_t* __restrict__ sc_cache =
        (kv == 0) ? key_scale_cache : value_scale_cache;

    // Source pointer for this token (use actual stride, not assumed contiguous)
    const CudaType* __restrict__ token_src =
        reinterpret_cast<const CudaType*>(src) + token_idx * src_stride;

    // Destination bases in data and scale caches for this token's block
    uint8_t* __restrict__ data_block =
        data_cache + block_idx * data_block_stride;
    uint8_t* __restrict__ scale_block =
        sc_cache + block_idx * scale_block_stride;

    for (int g = tg_id; g < total_groups; g += num_thread_groups) {
      const int head = g / groups_per_head;
      const int group_in_head = g % groups_per_head;

      // Load 16 (or 8) bf16 elements from source
      PVec in_vec;
      const CudaType* __restrict__ src_ptr =
          token_src + head * head_size + group_in_head * CVT_FP4_SF_VEC_SIZE +
          tg_lane * ELTS;

#pragma unroll
      for (int i = 0; i < ELTS / 2; i++) {
        in_vec.elts[i] = reinterpret_cast<
            const typename PackedTypeConverter<CudaType>::Type*>(src_ptr)[i];
      }

      // Quantize: produces packed fp4 and writes scale factor.
      uint8_t sf_val;
      uint8_t* sf_out_ptr = (tg_lane == 0) ? &sf_val : nullptr;

      fp4_packed_t packed = cvt_warp_fp16_to_fp4<CudaType, THREADS_PER_SF>(
          in_vec, global_scale, sf_out_ptr);

      // Write packed FP4 data to data cache
      uint8_t* __restrict__ data_dst = data_block + head * data_head_stride +
                                       block_offset * data_block_offset_stride;

#if CVT_FP4_PACK16
      {
        // 16 elements -> 8 bytes (u32x2)
        int data_byte_offset = group_in_head * 8;
        reinterpret_cast<uint64_t*>(data_dst + data_byte_offset)[0] =
            (uint64_t(packed.hi) << 32) | uint64_t(packed.lo);
      }
#else
      {
        // 8 elements -> 4 bytes (uint32_t)
        int data_byte_offset =
            group_in_head * CVT_FP4_SF_VEC_SIZE / 2 + tg_lane * ELTS / 2;
        reinterpret_cast<uint32_t*>(data_dst + data_byte_offset)[0] = packed;
      }
#endif

      // Write block scale to scale cache.
      // K (kv==0): always linear layout (no swizzle).
      // V (kv==1): swizzled layout for the SM100 trtllm-gen MHA kernel when
      //   swizzle_v_sf is true; linear (same as K) when false, which the
      //   FlashInfer FA2 paged nvfp4 reader on sm120/sm121 requires.
      if (sf_out_ptr != nullptr) {
        int scale_idx = group_in_head;
        uint8_t* __restrict__ scale_dst;
        if (kv == 0 || !swizzle_v_sf) {
          scale_dst = scale_block + head * scale_head_stride +
                      block_offset * scale_block_offset_stride + scale_idx;
        } else {
          int swizzled_offset =
              swizzle_scale_offset(block_offset, scale_idx, scale_dim);
          int swizzled_t = swizzled_offset / scale_dim;
          int swizzled_s = swizzled_offset % scale_dim;
          scale_dst = scale_block + head * scale_head_stride +
                      swizzled_t * scale_block_offset_stride + swizzled_s;
        }
        *scale_dst = sf_val;
      }
    }
  }
}

}  // namespace vllm

// Plain-ABI dispatch (the upstream dispatch lives in _C_stable_libtorch with
// the torch stable ABI; this standalone copy uses the standard ABI so it can
// be built with torch.utils.cpp_extension against the stock torch wheel).
void reshape_and_cache_nvfp4_standalone(torch::Tensor key, torch::Tensor value,
                                        torch::Tensor key_cache,
                                        torch::Tensor value_cache,
                                        torch::Tensor slot_mapping,
                                        torch::Tensor k_scale,
                                        torch::Tensor v_scale) {
  const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
  int num_tokens = slot_mapping.size(0);
  int num_heads = key.size(1);
  int head_size = key.size(2);
  int data_dim = head_size / 2;
  int scale_dim = head_size / 16;
  int full_dim = data_dim + scale_dim;

  // key_cache is kv_cache[:, 0] with shape
  // [num_blocks, block_size, num_heads, full_dim] in logical order.
  // Strides encode the physical layout (HND or NHD).
  TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D");
  TORCH_CHECK(key_cache.size(3) == full_dim,
              "key_cache last dim must be data_dim + scale_dim, got ",
              key_cache.size(3), " expected ", full_dim);

  int block_size = key_cache.size(1);

  TORCH_CHECK(head_size % 16 == 0,
              "head_size must be divisible by 16 for NVFP4 KV cache");

  // SM120/SM121 (consumer Blackwell) serve NVFP4 KV via the FlashInfer FA2
  // paged reader, which takes the scale-factor strides from the SF tensor
  // itself and reads V scales linearly; the SM100 trtllm-gen reader keeps
  // its 4-token V scale swizzle. Both consume the same per-page
  // [K_data | K_scale | V_data | V_scale] layout, so only the V-scale
  // swizzle is arch-conditional.
  const bool swizzle_v_sf = at::cuda::getCurrentDeviceProperties()->major < 12;

  TORCH_CHECK(!swizzle_v_sf || block_size % 4 == 0,
              "block_size must be divisible by 4 for NVFP4 KV cache V "
              "scale-factor swizzle (SM100 trtllm-gen path)");

  // Detect physical layout from strides (based on full_dim).
  // HND: head stride > block_offset stride.
  bool is_hnd = key_cache.stride(2) > key_cache.stride(1);

  int64_t data_block_stride = key_cache.stride(0);  // page_bytes
  int64_t data_head_stride, data_block_offset_stride;
  if (is_hnd) {
    data_head_stride = (int64_t)block_size * data_dim;
    data_block_offset_stride = data_dim;
  } else {
    data_head_stride = data_dim;
    data_block_offset_stride = (int64_t)num_heads * data_dim;
  }

  // Page layout: [K_data | K_scale | V_data | V_scale]
  // Scale follows data within each KV side.
  int64_t data_per_kv = (int64_t)num_heads * block_size * data_dim;

  uint8_t* key_scale_ptr = key_cache.data_ptr<uint8_t>() + data_per_kv;
  uint8_t* value_scale_ptr =
      value_cache.data_ptr<uint8_t>() + data_per_kv;

  // Scale strides: same page stride, inner strides from layout.
  int64_t scale_block_stride = data_block_stride;
  int64_t scale_head_stride, scale_block_offset_stride;
  if (is_hnd) {
    scale_head_stride = (int64_t)block_size * scale_dim;
    scale_block_offset_stride = scale_dim;
  } else {
    scale_head_stride = scale_dim;
    scale_block_offset_stride = (int64_t)num_heads * scale_dim;
  }

  const float* k_scale_ptr = k_scale.const_data_ptr<float>();
  const float* v_scale_ptr = v_scale.const_data_ptr<float>();

  int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE;
  int total_groups = num_heads * groups_per_head;
  constexpr int THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
  int num_threads = std::min(total_groups * THREADS_PER_SF, 512);
  num_threads = ((num_threads + 31) / 32) * 32;

  dim3 grid(num_tokens);
  dim3 block(num_threads);

  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  auto launch = [&](auto scalar) {
    using scalar_t = decltype(scalar);
    vllm::reshape_and_cache_nvfp4_kernel<scalar_t><<<grid, block, 0, stream>>>(
        key.const_data_ptr<scalar_t>(), value.const_data_ptr<scalar_t>(),
        key_cache.data_ptr<uint8_t>(), value_cache.data_ptr<uint8_t>(),
        key_scale_ptr, value_scale_ptr, slot_mapping.const_data_ptr<int64_t>(),
        k_scale_ptr, v_scale_ptr, key.stride(0), value.stride(0), num_heads,
        head_size, block_size, data_block_stride, data_head_stride,
        data_block_offset_stride, scale_block_stride, scale_head_stride,
        scale_block_offset_stride, swizzle_v_sf);
  };
  if (key.scalar_type() == at::kHalf) {
    launch(c10::Half());
  } else if (key.scalar_type() == at::kBFloat16) {
    launch(c10::BFloat16());
  } else {
    TORCH_CHECK(false, "reshape_and_cache_nvfp4 only supports fp16/bf16");
  }
}

TORCH_LIBRARY(_my_nvfp4_kv, m) {
  m.def(
      "reshape_and_cache_nvfp4(Tensor key, Tensor value, Tensor key_cache, "
      "Tensor value_cache, Tensor slot_mapping, Tensor k_scale, "
      "Tensor v_scale) -> ()");
}

TORCH_LIBRARY_IMPL(_my_nvfp4_kv, CUDA, m) {
  m.impl("reshape_and_cache_nvfp4", &reshape_and_cache_nvfp4_standalone);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("reshape_and_cache_nvfp4", &reshape_and_cache_nvfp4_standalone,
        py::arg("key"), py::arg("value"), py::arg("key_cache"),
        py::arg("value_cache"), py::arg("slot_mapping"), py::arg("k_scale"),
        py::arg("v_scale"));
}
CU_EOF

cat > "$SRC/nvfp4_utils.cuh" <<'H1_EOF'
/*
 * Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <cuda_runtime.h>
#include <cuda_fp8.h>
#include <utility>

#include "cuda_vec_utils.cuh"

#if defined(NVFP4_ENABLE_ELTS16) && defined(CUDART_VERSION) && \
    CUDART_VERSION >= 12090
  #define ELTS_PER_THREAD 16
  #define CVT_FP4_PACK16 1
constexpr int CVT_FP4_ELTS_PER_THREAD = 16;
#else
  #define ELTS_PER_THREAD 8
  #define CVT_FP4_PACK16 0
constexpr int CVT_FP4_ELTS_PER_THREAD = 8;
#endif

constexpr int CVT_FP4_SF_VEC_SIZE = 16;

namespace vllm {

// Explicit half/bf16 -> float conversion (CUDA 13 has no implicit conversion).
__device__ __forceinline__ float nvfp4_half_to_float(__half v) {
  return __half2float(v);
}
__device__ __forceinline__ float nvfp4_half_to_float(__nv_bfloat16 v) {
  return __bfloat162float(v);
}

template <typename Int>
__host__ __device__ inline Int round_up(Int x, Int y) {
  static_assert(std::is_integral_v<Int>,
                "round_up argument must be integral type");
  return ((x + y - 1) / y) * y;
}

template <typename Int>
__host__ __device__ __forceinline__ Int div_round_up(Int x, Int y) {
  return (x + y - 1) / y;
}

// Compute effective rows for grid configuration with swizzled SF layouts.
inline int computeEffectiveRows(int m) {
  constexpr int ROW_TILE = 128;
  return round_up(m, ROW_TILE);
}

// Compute the shape of the swizzled SF output tensor.
// Returns (rounded_m, rounded_n / 4) where:
//   rounded_m = round_up(m, 128)
//   rounded_n = round_up(n / CVT_FP4_SF_VEC_SIZE, 4)
inline std::pair<int64_t, int64_t> computeSwizzledSFShape(int64_t m,
                                                          int64_t n) {
  int64_t rounded_m = round_up(m, static_cast<int64_t>(128));
  int64_t scale_n = n / CVT_FP4_SF_VEC_SIZE;
  int64_t rounded_n = round_up(scale_n, static_cast<int64_t>(4));
  return {rounded_m, rounded_n / 4};
}

// Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t).
inline __device__ uint32_t fp32_vec8_to_e2m1(float (&array)[8]) {
  uint32_t val;
  asm volatile(
      "{\n"
      ".reg .b8 byte0;\n"
      ".reg .b8 byte1;\n"
      ".reg .b8 byte2;\n"
      ".reg .b8 byte3;\n"
      "cvt.rn.satfinite.e2m1x2.f32   byte0, %2, %1;\n"
      "cvt.rn.satfinite.e2m1x2.f32   byte1, %4, %3;\n"
      "cvt.rn.satfinite.e2m1x2.f32   byte2, %6, %5;\n"
      "cvt.rn.satfinite.e2m1x2.f32   byte3, %8, %7;\n"
      "mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
      "}"
      : "=r"(val)
      : "f"(array[0]), "f"(array[1]), "f"(array[2]), "f"(array[3]),
        "f"(array[4]), "f"(array[5]), "f"(array[6]), "f"(array[7]));
  return val;
}

// Convert 4 float2 values into 8 e2m1 values (represented as one uint32_t).
__device__ __forceinline__ uint32_t fp32_vec8_to_e2m1(float2 (&array)[4]) {
  uint32_t val;
  asm volatile(
      "{\n"
      ".reg .b8 byte0;\n"
      ".reg .b8 byte1;\n"
      ".reg .b8 byte2;\n"
      ".reg .b8 byte3;\n"
      "cvt.rn.satfinite.e2m1x2.f32   byte0, %2, %1;\n"
      "cvt.rn.satfinite.e2m1x2.f32   byte1, %4, %3;\n"
      "cvt.rn.satfinite.e2m1x2.f32   byte2, %6, %5;\n"
      "cvt.rn.satfinite.e2m1x2.f32   byte3, %8, %7;\n"
      "mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
      "}\n"
      : "=r"(val)
      : "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y),
        "f"(array[2].x), "f"(array[2].y), "f"(array[3].x), "f"(array[3].y));
  return val;
}

struct u32x2 {
  uint32_t lo, hi;
};

using fp4_packed_t = std::conditional_t<CVT_FP4_PACK16, u32x2, uint32_t>;

__device__ __forceinline__ u32x2 fp32_vec16_to_e2m1(float2 (&array)[8]) {
  u32x2 out;
  asm volatile(
      "{\n"
      ".reg .b8 b0;\n"
      ".reg .b8 b1;\n"
      ".reg .b8 b2;\n"
      ".reg .b8 b3;\n"
      ".reg .b8 b4;\n"
      ".reg .b8 b5;\n"
      ".reg .b8 b6;\n"
      ".reg .b8 b7;\n"
      "cvt.rn.satfinite.e2m1x2.f32   b0,  %3,  %2;\n"
      "cvt.rn.satfinite.e2m1x2.f32   b1,  %5,  %4;\n"
      "cvt.rn.satfinite.e2m1x2.f32   b2,  %7,  %6;\n"
      "cvt.rn.satfinite.e2m1x2.f32   b3,  %9,  %8;\n"
      "cvt.rn.satfinite.e2m1x2.f32   b4, %11, %10;\n"
      "cvt.rn.satfinite.e2m1x2.f32   b5, %13, %12;\n"
      "cvt.rn.satfinite.e2m1x2.f32   b6, %15, %14;\n"
      "cvt.rn.satfinite.e2m1x2.f32   b7, %17, %16;\n"
      "mov.b32 %0, {b0, b1, b2, b3};\n"
      "mov.b32 %1, {b4, b5, b6, b7};\n"
      "}\n"
      : "=r"(out.lo), "=r"(out.hi)
      : "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y),
        "f"(array[2].x), "f"(array[2].y), "f"(array[3].x), "f"(array[3].y),
        "f"(array[4].x), "f"(array[4].y), "f"(array[5].x), "f"(array[5].y),
        "f"(array[6].x), "f"(array[6].y), "f"(array[7].x), "f"(array[7].y));
  return out;
}

__device__ __forceinline__ uint32_t pack_fp4(float2 (&v)[4]) {
  return fp32_vec8_to_e2m1(v);
}

__device__ __forceinline__ u32x2 pack_fp4(float2 (&v)[8]) {
  return fp32_vec16_to_e2m1(v);
}

// Fast reciprocal.
__device__ __forceinline__ float reciprocal_approximate_ftz(float a) {
  float b;
  asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(b) : "f"(a));
  return b;
}

// Compute SF output offset for swizzled tensor core layout.
// SF layout: [numMTiles, numKTiles, 32, 4, 4]
// Caller must precompute: numKTiles = (numCols + 63) / 64
template <class SFType, int CVT_FP4_NUM_THREADS_PER_SF>
__device__ __forceinline__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset(
    int rowIdx, int colIdx, int32_t numKTiles, SFType* SFout) {
  static_assert(CVT_FP4_NUM_THREADS_PER_SF == 1 ||
                CVT_FP4_NUM_THREADS_PER_SF == 2);

  // One pair of threads write one SF to global memory.
  // TODO: stage through smem for packed STG.32
  // is it better than STG.8 from 4 threads ?
  if (threadIdx.x % CVT_FP4_NUM_THREADS_PER_SF != 0) {
    return nullptr;
  }

  // SF vector index (16 elements share one SF in the K dimension).
  int32_t kIdx = colIdx / CVT_FP4_NUM_THREADS_PER_SF;
  int32_t mIdx = rowIdx;

  // Decompose indices using bitwise ops (all divisors are powers of 2).
  // SF layout [numMTiles, numKTiles, 32 (mTile), 4 (mTile), 4(kTile)]
  int32_t mTileIdx = mIdx >> 7;         // mIdx / 128
  int32_t outerMIdx = mIdx & 31;        // mIdx % 32
  int32_t innerMIdx = (mIdx >> 5) & 3;  // (mIdx / 32) % 4
  int32_t kTileIdx = kIdx >> 2;         // kIdx / 4
  int32_t innerKIdx = kIdx & 3;         // kIdx % 4

  // Compute global SF offset: mTileIdx * (numKTiles * 512) + kTileIdx * 512 +
  //                           outerMIdx * 16 + innerMIdx * 4 + innerKIdx
  // Use bitwise OR for non-overlapping lower bits.
  int64_t SFOffset = (static_cast<int64_t>(mTileIdx) * numKTiles + kTileIdx)
                         << 9 |
                     (outerMIdx << 4) | (innerMIdx << 2) | innerKIdx;

  return reinterpret_cast<uint8_t*>(SFout) + SFOffset;
}

template <class SFType>
__device__ __forceinline__ uint8_t* sf_out_rowmajor_u8(int row, int pack,
                                                       int packs_per_row_sf,
                                                       SFType* SFout) {
  constexpr int PACK = CVT_FP4_ELTS_PER_THREAD;
  constexpr int THREADS_PER_SF =
      CVT_FP4_SF_VEC_SIZE / PACK;  // 1 if PACK=16, 2 else PACK=8

  if (threadIdx.x % THREADS_PER_SF != 0) return nullptr;

  int sf_col =
      pack / THREADS_PER_SF;  // PACK=16 => sf_col=pack; PACK=8 => sf_col=pack/2
  int64_t off = (int64_t)row * packs_per_row_sf + sf_col;

  return (uint8_t*)SFout + off;
}

// Quantizes the provided PackedVec into the uint32_t output
template <class Type, int CVT_FP4_NUM_THREADS_PER_SF, bool UE8M0_SF = false>
__device__ __forceinline__ fp4_packed_t cvt_warp_fp16_to_fp4(
    PackedVec<Type, CVT_FP4_PACK16>& vec, float SFScaleVal, uint8_t* SFout) {
  // Get absolute maximum values among the local 8 values.
  auto localMax = __habs2(vec.elts[0]);

  // Local maximum value.
#pragma unroll
  for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
    localMax = __hmax2(localMax, __habs2(vec.elts[i]));
  }

  // Get the absolute maximum among all 16 values (two threads).

  if constexpr (CVT_FP4_NUM_THREADS_PER_SF == 2) {
    localMax = __hmax2(__shfl_xor_sync(0xffffffffu, localMax, 1), localMax);
  }
  // Get the final absolute maximum values.
  // (explicit per-type conversion: CUDA 13 has no implicit __half/bf16->float)
  float vecMax = nvfp4_half_to_float(localMax.x);
  float localMaxY = nvfp4_half_to_float(localMax.y);
  vecMax = fmaxf(vecMax, localMaxY);

  // 8 bits representation of the SF.
  float SFValue;
  uint8_t fp8SFVal;

  if constexpr (UE8M0_SF) {
    // OCP MX spec E8M0 scale computation (MXFP4 path):
    // scale_exp = biased_exponent(round_up(vecMax)) - 2
    //   -2 because max E2M1 value is 6.0 ≈ 2^2.58; we use 2^2=4 as the
    //   safe divisor so that max_val / scale <= 6.0 for values near 2^n.
    uint32_t max_bits = __float_as_uint(vecMax);
    // Add rounding bias at mantissa bit 21 (equivalent to bf16 val_to_add=32
    // at bit 5). Threshold: values with mantissa >= 0.75 (i.e. >= 1.75*2^n)
    // round up to the next power of 2.
    uint32_t rounded_bits = (max_bits + (1u << 21)) & 0xFF800000u;
    uint32_t biased_exp = (rounded_bits >> 23) & 0xFFu;
    uint32_t scale_exp = (biased_exp > 2u) ? (biased_exp - 2u) : 0u;
    scale_exp = min(scale_exp, 254u);
    fp8SFVal = static_cast<uint8_t>(scale_exp);
    // Reconstruct scale as float32: scale = 2^(scale_exp - 127)
    uint32_t sf_bits = scale_exp << 23;
    SFValue = __uint_as_float(sf_bits);
  } else {
    // NVFP4 path: scale = max / 6.0, stored as E4M3.
    SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f));
    // Here SFValue is always positive, so E4M3 is the same as UE4M3.
    __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue);
    reinterpret_cast<__nv_fp8_e4m3&>(fp8SFVal) = tmp;
    // Convert back to fp32.
    SFValue = float(tmp);
  }

  // Write the SF to global memory (STG.8).
  if (SFout) *SFout = fp8SFVal;

  // Get the output scale (= 1 / SFValue for the MXFP4/UE8M0 path where
  // SFScaleVal=1).  Use exact division for UE8M0 to ensure bit-exact scaling
  // that matches the reference QDQ implementation (dividing by a power-of-2
  // scale is exact in IEEE 754).
  float outputScale;
  if constexpr (UE8M0_SF) {
    // SFValue is always a power of 2 for UE8M0, so 1/SFValue is exact.
    outputScale = SFValue != 0.0f ? (1.0f / SFValue) : 0.0f;
  } else {
    // NVFP4 path: use fast approximate reciprocal (original behavior).
    outputScale = SFValue != 0.0f
                      ? reciprocal_approximate_ftz(
                            SFValue * reciprocal_approximate_ftz(SFScaleVal))
                      : 0.0f;
  }

  // Convert the input to float.
  float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];

#pragma unroll
  for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
    fp2Vals[i] = cast_to_float2(vec.elts[i]);
    fp2Vals[i].x *= outputScale;
    fp2Vals[i].y *= outputScale;
  }

  // Convert to e2m1 values.
  return pack_fp4(fp2Vals);
}

// silu in float32
__device__ __forceinline__ float silu(float x) {
  return __fdividef(x, (1.f + __expf(-x)));
}

__device__ __forceinline__ float2 silu2(float2 x) {
  return make_float2(silu(x.x), silu(x.y));
}

template <class Type>
__inline__ __device__ PackedVec<Type, CVT_FP4_PACK16> compute_silu_mul(
    const PackedVec<Type, CVT_FP4_PACK16>& x_vec,
    const PackedVec<Type, CVT_FP4_PACK16>& y_vec) {
  PackedVec<Type, CVT_FP4_PACK16> result;

#pragma unroll
  for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; ++i) {
    // silu_mul in float32
    using packed_t = typename PackedTypeConverter<Type>::Type;
    float2 silu_vec = silu2(cast_to_float2(x_vec.elts[i]));
    float2 y_f2 = cast_to_float2(y_vec.elts[i]);
    result.elts[i] = cast_to_packed<packed_t>(
        make_float2(silu_vec.x * y_f2.x, silu_vec.y * y_f2.y));
  }
  return result;
}

}  // namespace vllm
H1_EOF

cat > "$SRC/cuda_vec_utils.cuh" <<'H2_EOF'
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project

#pragma once

#include <torch/headeronly/util/BFloat16.h>
#include <torch/headeronly/util/Half.h>
#include <cassert>

#ifdef USE_ROCM
  #include <hip/hip_runtime.h>
  #include <hip/hip_bf16.h>
  #include <hip/hip_fp16.h>
#else
  #include <cuda_bf16.h>
  #include <cuda_fp16.h>
  #include <cuda_runtime.h>
#endif

// Device-side: SM100+ architecture with CUDA 12.9+ toolkit, which
// together enable 256-bit (v8.u32) PTX load/store instructions.
// Use for PTX instruction selection with architecture fallback paths.
#if !defined(USE_ROCM) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \
    defined(CUDART_VERSION) && CUDART_VERSION >= 12090
  #define VLLM_256B_PTX_ENABLED 1
#else
  #define VLLM_256B_PTX_ENABLED 0
#endif

namespace vllm {

// ============================================================
// Types and traits
// ============================================================

// 256-bit (32-byte) aligned vector type: 8 x uint32_t
struct alignas(32) u32x8_t {
  uint32_t d[8];
};

// VecTraits: select between 128-bit (int4) and 256-bit
// (u32x8_t) vector types at compile time.
template <bool support_256>
struct VecTraits;

template <>
struct VecTraits<true> {
  static constexpr int ARCH_MAX_VEC_SIZE = 32;
  using vec_t = u32x8_t;
};

template <>
struct VecTraits<false> {
  static constexpr int ARCH_MAX_VEC_SIZE = 16;
  using vec_t = int4;
};

// PackedTypeConverter: map between CUDA scalar and packed types
//   half  <-> half2,  __nv_bfloat16 <-> __nv_bfloat162, etc.
template <typename T>
struct PackedTypeConverter {
  static_assert(sizeof(T) == 0,
                "PackedTypeConverter is not specialized for this type.");
};

template <>
struct PackedTypeConverter<half2> {
  using Type = half;
};

template <>
struct PackedTypeConverter<half> {
  using Type = half2;
};

template <>
struct PackedTypeConverter<__nv_bfloat162> {
  using Type = __nv_bfloat16;
};

template <>
struct PackedTypeConverter<__nv_bfloat16> {
  using Type = __nv_bfloat162;
};

template <>
struct PackedTypeConverter<float> {
  using Type = float2;
};

template <>
struct PackedTypeConverter<float2> {
  using Type = float;
};

template <>
struct PackedTypeConverter<c10::Half> {
  using Type = half2;
};

template <>
struct PackedTypeConverter<c10::BFloat16> {
  using Type = __nv_bfloat162;
};

// CUDATypeConverter: map PyTorch scalar types to CUDA scalar
//   c10::Half -> half,  c10::BFloat16 -> __nv_bfloat16
template <typename T>
struct CUDATypeConverter {
  using Type = T;
};

template <>
struct CUDATypeConverter<c10::Half> {
  using Type = half;
};

template <>
struct CUDATypeConverter<c10::BFloat16> {
  using Type = __nv_bfloat16;
};

// PackedVec: typed vector container for packed element access.
//   Derives alignment and element count from VecTraits.
//   Type is the CUDA scalar type (e.g. half, __nv_bfloat16).
template <class Type, bool use_256b>
struct alignas(VecTraits<use_256b>::ARCH_MAX_VEC_SIZE) PackedVec {
  static constexpr int NUM_ELTS =
      VecTraits<use_256b>::ARCH_MAX_VEC_SIZE /
      sizeof(typename PackedTypeConverter<Type>::Type);
  typename PackedTypeConverter<Type>::Type elts[NUM_ELTS];
};

// ============================================================
// Load / store primitives
// ============================================================

// 256-bit load / store: SM100+ only (PTX v8 instructions).
__device__ __forceinline__ void ld256(u32x8_t& val, const u32x8_t* ptr) {
#if VLLM_256B_PTX_ENABLED
  asm volatile("ld.global.nc.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8];\n"
               : "=r"(val.d[0]), "=r"(val.d[1]), "=r"(val.d[2]), "=r"(val.d[3]),
                 "=r"(val.d[4]), "=r"(val.d[5]), "=r"(val.d[6]), "=r"(val.d[7])
               : "l"(ptr));
#else
  assert(false && "ld256 requires SM100+ with CUDA 12.9+");
#endif
}

__device__ __forceinline__ void st256(u32x8_t& val, u32x8_t* ptr) {
#if VLLM_256B_PTX_ENABLED
  asm volatile("st.global.v8.u32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};\n"
               :
               : "l"(ptr), "r"(val.d[0]), "r"(val.d[1]), "r"(val.d[2]),
                 "r"(val.d[3]), "r"(val.d[4]), "r"(val.d[5]), "r"(val.d[6]),
                 "r"(val.d[7])
               : "memory");
#else
  assert(false && "st256 requires SM100+ with CUDA 12.9+");
#endif
}

// Generic ld256 / st256 for any 32-byte aligned type (e.g. PackedVec).
// Non-template overloads above are preferred for u32x8_t.
template <typename T>
__device__ __forceinline__ void ld256(T& val, const T* ptr) {
  static_assert(sizeof(T) == 32, "ld256 requires a 32-byte type");
  ld256(reinterpret_cast<u32x8_t&>(val), reinterpret_cast<const u32x8_t*>(ptr));
}

template <typename T>
__device__ __forceinline__ void st256(T& val, T* ptr) {
  static_assert(sizeof(T) == 32, "st256 requires a 32-byte type");
  st256(reinterpret_cast<u32x8_t&>(val), reinterpret_cast<u32x8_t*>(ptr));
}

// 128-bit load / store via __ldg (read-only cache hint).
template <typename T>
__device__ __forceinline__ void ld128(T& val, const T* ptr) {
  static_assert(sizeof(T) == 16, "ld128 requires a 16-byte type");
  *reinterpret_cast<int4*>(&val) = __ldg(reinterpret_cast<const int4*>(ptr));
}

template <typename T>
__device__ __forceinline__ void st128(T& val, T* ptr) {
  static_assert(sizeof(T) == 16, "st128 requires a 16-byte type");
  *reinterpret_cast<int4*>(ptr) = *reinterpret_cast<int4*>(&val);
}

// 256-bit cache-streaming (.cs) load / store: SM100+ only.
__forceinline__ __device__ u32x8_t ld256_cs(const u32x8_t* addr) {
#if VLLM_256B_PTX_ENABLED
  u32x8_t val;
  asm volatile("ld.global.cs.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8];"
               : "=r"(val.d[0]), "=r"(val.d[1]), "=r"(val.d[2]), "=r"(val.d[3]),
                 "=r"(val.d[4]), "=r"(val.d[5]), "=r"(val.d[6]), "=r"(val.d[7])
               : "l"(addr));
  return val;
#else
  assert(false && "ld256_cs requires SM100+ with CUDA 12.9+");
  return u32x8_t{};
#endif
}

__forceinline__ __device__ void st256_cs(u32x8_t* addr, u32x8_t val) {
#if VLLM_256B_PTX_ENABLED
  asm volatile(
      "st.global.cs.v8.u32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" ::"l"(addr),
      "r"(val.d[0]), "r"(val.d[1]), "r"(val.d[2]), "r"(val.d[3]), "r"(val.d[4]),
      "r"(val.d[5]), "r"(val.d[6]), "r"(val.d[7]));
#else
  assert(false && "st256_cs requires SM100+ with CUDA 12.9+");
#endif
}

// 32-bit load / store.
__device__ __forceinline__ int ld32(const int* addr) { return __ldg(addr); }

__device__ __forceinline__ void st32(int* addr, int val) { *addr = val; }

// 32-bit cache-streaming (.cs) load / store.
// Falls back to ld32/st32 on ROCm (no .cs hint).
__forceinline__ __device__ int ld32_cs(const int* addr) {
  int val;
#ifndef USE_ROCM
  asm volatile("ld.global.cs.b32 %0, [%1];" : "=r"(val) : "l"(addr));
#else
  val = ld32(addr);
#endif
  return val;
}

__forceinline__ __device__ void st32_cs(int* addr, int val) {
#ifndef USE_ROCM
  asm volatile("st.global.cs.b32 [%0], %1;" ::"l"(addr), "r"(val));
#else
  st32(addr, val);
#endif
}

// 128-bit cache-streaming (.cs) load / store.
// Falls back to ld128/st128 on ROCm (no .cs hint).
__forceinline__ __device__ int4 ld128_cs(const int4* addr) {
  int4 val;
#ifndef USE_ROCM
  asm volatile("ld.global.cs.v4.u32 {%0,%1,%2,%3}, [%4];"
               : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w)
               : "l"(addr));
#else
  ld128(val, addr);
#endif
  return val;
}

__forceinline__ __device__ void st128_cs(int4* addr, int4 val) {
#ifndef USE_ROCM
  asm volatile("st.global.cs.v4.u32 [%0], {%1,%2,%3,%4};" ::"l"(addr),
               "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w));
#else
  st128(val, addr);
#endif
}

// Predicated 256-bit / 128-bit cache-global (.cg) loads.
// Returns zero if pred is false.  SM100+ only.
__device__ __forceinline__ void ld256_cg_or_zero(u32x8_t& val, const void* ptr,
                                                 bool pred) {
#if VLLM_256B_PTX_ENABLED
  asm volatile(
      "{\n"
      "  .reg .pred pr;\n"
      "  setp.ne.u32 pr, %8, 0;\n"
      "  mov.u32 %0, 0;\n"
      "  mov.u32 %1, 0;\n"
      "  mov.u32 %2, 0;\n"
      "  mov.u32 %3, 0;\n"
      "  mov.u32 %4, 0;\n"
      "  mov.u32 %5, 0;\n"
      "  mov.u32 %6, 0;\n"
      "  mov.u32 %7, 0;\n"
      "  @pr ld.global.cg.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%9];\n"
      "}\n"
      : "=r"(val.d[0]), "=r"(val.d[1]), "=r"(val.d[2]), "=r"(val.d[3]),
        "=r"(val.d[4]), "=r"(val.d[5]), "=r"(val.d[6]), "=r"(val.d[7])
      : "r"((int)pred), "l"(ptr));
#else
  assert(false && "ld256_cg_or_zero requires SM100+ with CUDA 12.9+");
#endif
}

__device__ __forceinline__ void ld128_cg_or_zero(uint4& val, const void* ptr,
                                                 bool pred) {
#ifndef USE_ROCM
  uint32_t r0, r1, r2, r3;

  asm volatile(
      "{\n"
      "  .reg .pred pr;\n"
      "  setp.ne.u32 pr, %4, 0;\n"
      "  mov.u32 %0, 0;\n"
      "  mov.u32 %1, 0;\n"
      "  mov.u32 %2, 0;\n"
      "  mov.u32 %3, 0;\n"
      "  @pr ld.global.cg.v4.u32 {%0,%1,%2,%3}, [%5];\n"
      "}\n"
      : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
      : "r"((int)pred), "l"(ptr));

  val = uint4{r0, r1, r2, r3};
#else
  assert(false && "ld128_cg_or_zero is not supported on ROCm");
#endif
}

// ============================================================
// Alignment helpers
// ============================================================

__host__ __device__ __forceinline__ bool is_16byte_aligned(const void* ptr) {
  return (reinterpret_cast<uintptr_t>(ptr) & 15) == 0;
}

__host__ __device__ __forceinline__ bool is_32byte_aligned(const void* ptr) {
  return (reinterpret_cast<uintptr_t>(ptr) & 31) == 0;
}

// ============================================================
// Packed type conversion and arithmetic
// ============================================================

template <typename packed_t>
__device__ __forceinline__ float2 cast_to_float2(const packed_t& val) {
  if constexpr (std::is_same_v<packed_t, __nv_bfloat162>) {
    return __bfloat1622float2(val);
  } else if constexpr (std::is_same_v<packed_t, __half2>) {
    return __half22float2(val);
  } else if constexpr (std::is_same_v<packed_t, float2>) {
    return float2(val);
  }
}

template <typename packed_t>
__device__ __forceinline__ packed_t cast_to_packed(const float2& val) {
  if constexpr (std::is_same_v<packed_t, __nv_bfloat162>) {
    return __float22bfloat162_rn(val);
  } else if constexpr (std::is_same_v<packed_t, __half2>) {
    return __float22half2_rn(val);
  } else if constexpr (std::is_same_v<packed_t, float2>) {
    return float2(val);
  }
}

template <typename packed_t>
__device__ __forceinline__ packed_t packed_mul(const packed_t& x,
                                               const packed_t& y) {
  if constexpr (std::is_same_v<packed_t, __nv_bfloat162> ||
                std::is_same_v<packed_t, __half2>) {
    return __hmul2(x, y);
  } else if constexpr (std::is_same_v<packed_t, float2>) {
    return make_float2(x.x * y.x, x.y * y.y);
  }
}

}  // namespace vllm
H2_EOF

cat > "$SRC/build_ext.py" <<'PY_EOF'
import os

cuda_home = "/usr/local/lib/python3.12/dist-packages/nvidia/cu13"
os.environ["CUDA_HOME"] = cuda_home
os.environ["CUDA_PATH"] = cuda_home
# The pip cu13 wheel ships versioned .so.13 libs but no unversioned symlinks;
# the linker needs libcudart.so etc. Create them (the L8 landmine).
import glob

_libdir = f"{cuda_home}/lib"
for _f in glob.glob(os.path.join(_libdir, "*.so.*")):
    _base = _f.split(".so.")[0] + ".so"
    if not os.path.exists(_base):
        os.symlink(os.path.basename(_f), _base)
os.environ["TORCH_CUDA_ARCH_LIST"] = "12.0"

# vLLM's CMake builds Blackwell with family-conditional targets (12.0f) for
# fp4 instructions; torch's cpp_extension does not parse the 'f' suffix, so
# monkeypatch the arch-flags helper to emit the sm_120f gencode directly.
import torch.utils.cpp_extension as _cpp

_orig_arch_flags = _cpp._get_cuda_arch_flags


def _my_arch_flags(cflags=None):
    return ["-gencode=arch=compute_120f,code=sm_120f"]


_cpp._get_cuda_arch_flags = _my_arch_flags
# Invalidate any cached flags (older torch caches the result).
try:
    _orig_arch_flags.cache_clear()
except AttributeError:
    pass
os.environ.setdefault(
    "LIBRARY_PATH", f"{cuda_home}/lib"
)
os.environ.setdefault(
    "LD_LIBRARY_PATH", f"{cuda_home}/lib"
)

from torch.utils.cpp_extension import load

ext = load(
    name="my_nvfp4_kv",
    sources=["/src/my_nvfp4_kv.cu"],
    extra_cuda_cflags=[
        "-O3",
        "-DCCCL_DISABLE_CTK_COMPATIBILITY_CHECK",
        f"-I{cuda_home}/include/cccl",
    ],
    verbose=True,
)
print("BUILT:", ext.__file__)
PY_EOF


# 2. ship them to the box
ssh "$BOX" "mkdir -p $BOX_SRC && rm -rf $BOX_SRC/*"
scp -q "$SRC"/* "$BOX:$BOX_SRC/"

# 3. compile the standalone writer kernel (sm_120f) inside the base image
ssh "$BOX" "sudo podman run --rm --entrypoint python3 \
  -v $BOX_SRC:/src:Z -v $BOX_SRC/cache:/root/.cache/torch_extensions:Z \
  docker.io/vllm/vllm-openai:nightly /src/build_ext.py | tail -1"

# 4. apply the flashinfer.py patch inside a throwaway container, produce the
#    patched file on the bind mount
ssh "$BOX" "sudo podman run --rm --entrypoint bash -v $BOX_SRC:/src:Z \
  docker.io/vllm/vllm-openai:nightly -c '
    set -e
    command -v patch >/dev/null 2>&1 || apt-get update -qq >/dev/null \
      && apt-get install -y -qq patch >/dev/null 2>&1
    python3 - <<PY
import re, pathlib
p = pathlib.Path("/src/flashinfer.py.diff")
t = p.read_text()
t = re.sub(r"^--- .*", "--- flashinfer.py", t, count=1, flags=re.M)
t = re.sub(r"^\+\+\+ .*", "+++ flashinfer.py", t, count=1, flags=re.M)
p.write_text(t)
PY
    cd /usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends
    patch -p0 < /src/flashinfer.py.diff
    cp flashinfer.py /src/flashinfer.py.patched
'"

# 5. assemble the patched image
ssh "$BOX" 'set -e
sudo podman rm -f vllm-nvfp4kv-tmp 2>/dev/null || true
sudo podman create --name vllm-nvfp4kv-tmp --entrypoint /bin/bash docker.io/vllm/vllm-openai:nightly sleep infinity >/dev/null
sudo podman cp '$BOX_SRC'/flashinfer.py.patched vllm-nvfp4kv-tmp:/usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/flashinfer.py
sudo podman cp '$BOX_SRC'/cache/py312_cu130/my_nvfp4_kv/my_nvfp4_kv.so vllm-nvfp4kv-tmp:/usr/local/lib/python3.12/dist-packages/my_nvfp4_kv.so
sudo podman commit --change "ENTRYPOINT [\"vllm\", \"serve\"]" --change "CMD []" -q vllm-nvfp4kv-tmp localhost/vllm-nvfp4kv:latest
sudo podman rm vllm-nvfp4kv-tmp >/dev/null
echo ">> done: localhost/vllm-nvfp4kv:latest ready"'

3. Power profile (LACT)

Optional but recommended: cap the card at 400 W and hold the clock range. The GPU is managed by lactd (install via flatpak, run headless as a system service). Write this profile to /etc/lact/config.yaml (adjust the GPU ID to your card):

gpus:
  10DE:2B85-1458:416E-0000:01:00.0:
    power_cap: 400
    min_core_clock: 300
    max_core_clock: 2500
    gpu_clock_offsets:
      0: 600

Then systemctl restart lactd; a restart is required, not just writing the file: replacing the file inode (install) is missed by lactd's file watcher. Confirm with nvidia-smi --query-gpu=power.limit,clocks.gr. A reboot resets to the stock profile; re-apply after every boot. No LACT? The nvidia-smi fallback is nvidia-smi -pl 400 && nvidia-smi -lgc 300,2500 (runtime-only, same reboot caveat).

4. Launch

Serve with the patched image (adjust the HF-cache volume path to your setup):

sudo podman run --name vllm-serve --rm \
  --device nvidia.com/gpu=0 \
  --security-opt label=disable \
  --shm-size 16g \
  -p 8000:8000 \
  -e VLLM_KV_CACHE_LAYOUT=HND \
  -v "$HOME/.cache/huggingface:/root/.cache/huggingface:ro" \
  localhost/vllm-nvfp4kv:latest \
  gittensor-model-hub/Qwen3.8-27B-NVFP4-RTX5090 \
    --served-model-name qwen3.8-27b \
    --host 0.0.0.0 --port 8000 \
    --max-model-len 196608 \
    --gpu-memory-utilization 0.98 \
    --max-num-seqs 4 \
    --kv-cache-dtype nvfp4 \
    --enable-prefix-caching \
    --enable-auto-tool-choice \
    --trust-remote-code \
    --reasoning-parser qwen3 \
    --tool-call-parser qwen3_coder \
    --default-chat-template-kwargs '{"preserve_thinking": true}' \
    --override-generation-config '{"temperature":1.0,"top_p":0.95,"top_k":20,"min_p":0.0,"presence_penalty":0.0,"repetition_penalty":1.0}' \
    --compilation-config '{"cudagraph_mode":"piecewise"}' \
    --speculative-config '{"method":"mtp","num_speculative_tokens":3}'

Notes:

  • VLLM_KV_CACHE_LAYOUT=HND must reach the container (-e NAME=VALUE, not a bare export); without it the nvfp4 layout resolves wrong and output is garbage.
  • --compilation-config '{"cudagraph_mode":"piecewise"}' is mandatory; do not override it (see the gotcha).
  • Server is OpenAI-compatible on port 8000, model name qwen3.8-27b.

5. Validation checklist

  1. Boot log must show: Using LBHNC KV cache layout, cudagraph_mode: PIECEWISE, Capturing CUDA graphs (mixed prefill-decode, PIECEWISE): 100%, and GPU KV cache size: 451,041 tokens.
  2. KV pool ratio is the reliable check: ~451K tokens at 196K ctx (≈2.29x). A silent fallback to fp8 keeps the old small pool.
  3. Greedy determinism: same prompt ×6 → byte-identical output (all should produce content, not reasoning-only).
  4. Reasoning sanity: a simple "count 1-5" must return content after brief reasoning (empty-content + reasoning-only = cudagraph corruption; see failure modes).
  5. Long-context needle: ~90K-token haystack, retrieve a planted word.
  6. Vision: an actual image request (the checkpoint is multimodal).
  7. Quality gate: uvx llama-benchy --base-url http://<box>:8000/v1 --skip-coherence --pp 4000 8000 --tg 64 → pp ~11.4K / ~10.5K, tg64 ~130 / ~150.

6. Failure modes

Symptom Cause Fix
--kv-cache-dtype nvfp4 requires the SM100 ... path stock image use the patched image (rebuild)
... requires the HND KV cache layout layout resolved to NHD VLLM_KV_CACHE_LAYOUT=HND into the container
boots, output is !!!!!!!!... NHD layout, or V-scale swizzle (stock writer) HND env; use the standalone writer
boots, output is reasoning loops / empty content / hallucinated text, no crash FULL cudagraphs + MTP + nvfp4 PIECEWISE (--compilation-config '{"cudagraph_mode":"piecewise"}') or --enforce-eager
CUDA error: unspecified launch failure on first request model not loadable on the nvfp4 path (e.g. FP8-attention ModelOpt checkpoints) that checkpoint is SGLang-first; not a config issue
Engine core init failure (early nightly) stale image mix rebuild from a clean nightly

7. Upstream status (2026-08-23)

  • vLLM PRs #46329 and #50288 implement SM120 NVFP4 KV upstream. Still open. When they merge, stock vLLM accepts --kv-cache-dtype nvfp4 on SM120 and this patch becomes obsolete.
  • Until then, localhost/vllm-nvfp4kv:latest is required; everything else is stock.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment