[mxfp8-kv] Skip writes to the reserved CUDA-graph padding slot (#35351)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Ke Bao <ispobaoke@gmail.com>
Co-authored-by: Sam Shleifer <sam@thinkingmachines.ai>

python/sglang/kernels/ops/quantization/mxfp8_interleave_sf.py

@@ -40,8 +40,9 @@ def _store_sf_interleaved_kernel(
     tok_offsets = tok_start + tl.arange(0, BLOCK_T)
     mask = tok_offsets < num_tokens
 
-    # Load slot indices
+    # Slot 0 is the reserved CUDA-graph padding sink; skip writes to it.
     slots = tl.load(loc_ptr + tok_offsets, mask=mask, other=0)
+    mask = mask & (slots != 0)
     page_offsets = slots % page_size
     page_idxs = slots // page_size
 

python/sglang/kernels/ops/quantization/mxfp8_quant.py

@@ -143,6 +143,9 @@ def _mxfp8_quant_store_qkv_kernel(
         tl.store(sfq_ptr + (t * NQ + r) * SF + blk, sf)
     else:
         myloc = tl.load(loc_ptr + t).to(tl.int64)
+        # Slot 0 is the reserved CUDA-graph padding sink; skip writes to it.
+        if myloc == 0:
+            return
         if r < NQ + NKV:
             h = r - NQ
             cache = kc_ptr

python/sglang/srt/mem_cache/memory_pool.py

@@ -3650,20 +3650,28 @@ class MHATokenToKVPoolMXFP8(MHATokenToKVPool):
             )
             return
 
-        from sglang.srt.model_executor.runner import get_is_capture_mode
-
-        if get_is_capture_mode() and self.alt_stream is not None:
-            current_stream = self.device_module.current_stream()
-            self.alt_stream.wait_stream(current_stream)
-            self.k_buffer[idx][loc] = cache_k
-            self._write_scales(idx, loc, k_scale, v_scale)
-            with self.device_module.stream(self.alt_stream):
-                self.v_buffer[idx][loc] = cache_v
-            current_stream.wait_stream(self.alt_stream)
-        else:
-            self.k_buffer[idx][loc] = cache_k
-            self.v_buffer[idx][loc] = cache_v
-            self._write_scales(idx, loc, k_scale, v_scale)
+        # store_cache and store_sf_interleaved skip the reserved CUDA-graph
+        # padding slot 0 in-kernel, matching the bf16 pool.
+        row_bytes = self.head_num * self.head_dim * self.store_dtype.itemsize
+        v_row_bytes = self.head_num * self.v_head_dim * self.store_dtype.itemsize
+        assert _is_cuda and can_use_store_cache(row_bytes, v_row_bytes), (
+            f"MXFP8 KV cache requires CUDA and store_cache-compatible rows, "
+            f"got _is_cuda={_is_cuda}, {row_bytes=}, {v_row_bytes=}"
+        )
+        assert self.mxfp8_sf_interleaved, (
+            "MXFP8 KV cache requires the page_size=128 interleaved scale layout"
+        )
+        store_cache(
+            cache_k.reshape(loc.shape[0], -1),
+            cache_v.reshape(loc.shape[0], -1),
+            self.k_buffer[idx].view(-1, row_bytes // self.store_dtype.itemsize),
+            self.v_buffer[idx].view(-1, v_row_bytes // self.store_dtype.itemsize),
+            loc,
+            row_bytes=row_bytes,
+            v_row_bytes=v_row_bytes,
+            size_limit=self.size + self.page_size,
+        )
+        self._write_scales(idx, loc, k_scale, v_scale)
 
     def _write_scales(self, idx, loc, k_scale, v_scale):
         """Write per-token UE8M0 K/V scales — interleaved into the FA4