[Quant] Reject non-finite and non-per-tensor fp8 K/V scales at checkpoint load

[rchalamala]

Motivation

BaseKVCacheMethod.process_weights_after_loading in python/sglang/srt/layers/quantization/kv_cache.py validates checkpoint-loaded fp8 K/V scales only incidentally, so corrupt scales either crash with confusing errors or are silently accepted:

• Multi-element (e.g. per-channel) scales never reach the isinstance(k_scale, float) per-tensor check at the bottom of the method: if layer.k_scale > 0.0 raises RuntimeError: Boolean value of Tensor with more than one value is ambiguous first.
• NaN satisfies neither > 0.0 nor <= 0.0, so it falls into the single-kv_scale remap branch and dies on assert layer.k_scale > 0.0 — a bare AssertionError with no message.
• +inf passes > 0.0 and is copied straight into layer.k_scale_float, so an infinite scale reaches the attention forward.
• -inf in one scale alongside a valid other scale takes the remap branch (max(k, v) picks the valid one) and the corrupt value is silently replaced.
• On FNUZ hardware the loaded scales are doubled after any load-time check, so a finite scale above half the float32 maximum (e.g. 3e38) passes a pre-conversion finite check and becomes inf when copied back into the float32 parameter.

Probe of process_weights_after_loading on unmodified base 929230a (CPU, scalar/tensor parameters injected into a stub layer):

multi-element [1.0, 2.0]: RuntimeError: Boolean value of Tensor with more than one value is ambiguous
nan k_scale: AssertionError:
+inf k_scale: ACCEPTED k_float=inf v_float=1.0
-inf v_scale (valid k): ACCEPTED k_float=1.0 v_float=1.0
fnuz 3e38 k_scale: ACCEPTED k_float=6.0000000109955115e+38 v_float=2.0

Same probe with this PR:

multi-element [1.0, 2.0]: ValueError: Only support per-tensor scaling factor for fp8 KV cache, got k_scale shape=(2,) v_scale shape=()
nan k_scale: ValueError: fp8 KV cache scales must be finite, got k_scale=nan v_scale=1.0
+inf k_scale: ValueError: fp8 KV cache scales must be finite, got k_scale=inf v_scale=1.0
-inf v_scale (valid k): ValueError: fp8 KV cache scales must be finite, got k_scale=1.0 v_scale=-inf
fnuz 3e38 k_scale: ValueError: fp8 KV cache scales must be finite and fit in torch.float32, got k_scale=6.0000000109955115e+38 v_scale=2.0

Modifications

Two stacked guards at the top of process_weights_after_loading, plus a post-conversion guard before the layer is mutated:

1. Reject non-per-tensor scales before any scalar comparison: if either scale has numel() != 1, raise a ValueError naming both shapes. This replaces the ambiguous-bool RuntimeError with an explicit message.
2. Reject non-finite scales: copy both scales to CPU and require torch.isfinite(...).all(), otherwise raise a ValueError naming both values. This catches NaN and ±inf in one place, including the -inf case that was previously remapped away.
3. Re-validate the final converted scales after the FNUZ doubling / default / remap branches: require math.isfinite on both Python floats and require each to fit in torch.finfo(layer.k_scale.dtype).max, otherwise raise a ValueError naming both converted values. This closes the FNUZ overflow hole where a finite checkpoint scale doubles past the float32 maximum.

The existing branches are unchanged: separate positive scales still win, two non-positive scales still default to 1.0, and the single-kv_scale → k_scale → v_scale remap keeps its current behavior. The pre-existing isinstance(float) check stays as a backstop.

Accuracy Tests

New regression file test/registered/unit/layers/quantization/test_kv_cache_scale_validation.py (CPU, registered base-a-test-cpu): inf/-inf/NaN rejection, multi-element rejection with a per-tensor ValueError (including a multi-element scale containing NaN, which previously raised the ambiguous RuntimeError), zero-scale AssertionError preserved, finite scales accepted, unset scales defaulting to 1.0, FNUZ doubling of 3e38 rejected as post-conversion overflow, and FNUZ doubling of 1e38 accepted within float32.

Commands (CPU):

PYTHONPATH=<worktree>/python \
  python -m pytest test/registered/unit/layers/quantization/test_kv_cache_scale_validation.py -v
# 9 passed in 6.82s

PYTHONPATH=<worktree>/python \
  python -m pytest test/registered/unit/layers/quantization/test_compressed_tensors_kv_cache.py
# 4 passed in 6.78s

PYTHONPATH=<worktree>/python \
  python -m pytest test/registered/unit/layers/quantization/test_fp4_kv_cache_quant_method.py
# 22 passed, 1 skipped, 2 failed in 12.30s

The 2 fp4 failures (TestFP4MXBlock16KVCacheMethod::test_quantize_dequantize_roundtrip_cpu, TestFP4MXBlock16KVQuantizeUtil::test_roundtrip_cpu) are pre-existing and environmental: torch inductor emits ImportError: /lib64/libc.so.6: version 'GLIBC_2.38' not found on this CPU box. Re-running the same file on unmodified base 929230a reproduces the identical 2 failed / 22 passed / 1 skipped result.

Lint: ruff check --select=F401,F821,UP037 (the pre-commit hook's selection) passes on both changed files; ruff format --check reports both files already formatted.

Speed Tests and Profiling

No speed impact claimed; not measured (CPU-only validation host).

Validation limits

• CPU-only, unit-level: no GPU run and no end-to-end model load with a real fp8 checkpoint in this PR.
• The FNUZ doubling path is exercised only through a mocked is_fp8_fnuz() in the new tests (post-conversion overflow rejected, within-float32 doubling accepted); no real FNUZ/ROCm hardware run.
• The 2 fp4 KV-cache test failures are environmental (GLIBC) and equally present on the base commit.

Checklist

• Format your code according to the Format code with pre-commit (ruff check/format clean on changed files).
• Add unit tests according to the Run and add unit tests (see Accuracy Tests).
• Update documentation (no user-facing docs changes in this PR).
• Provide accuracy and speed benchmark results (CPU-only validation here; see Validation limits).
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #35421236779
Latest PR Test (Extra): ❌ Run #35421236633
Latest PR Test (AMD ROCm 10): ❌ Run #35421236734

[gilfordting]

/tag-and-rerun-ci