flashinfer.mla.cake_dsv4_nvfp4_kv_rope_quantize_insert

flashinfer.mla.cake_dsv4_nvfp4_kv_rope_quantize_insert(kv: Tensor, cache: Tensor, slot_mapping: Tensor, positions: Tensor, cos_sin_cache: Tensor, *, compress_ratio: int = 1) → None

Rotate, quantize and insert latent KV rows into the DeepSeek-V4 NVFP4 cache (KV only).

The compressed-pool and speculative-context writer: row i is inserted when i < slot_mapping.numel(), its slot is valid and (positions[i] + 1) % compress_ratio == 0; the RoPE row is cos_sin_cache[positions[i] // compress_ratio * compress_ratio]. With compress_ratio=1 every addressed row is inserted at its own position. The cache bytes equal flashinfer.mla.nvfp4_quantize_append_sparse_mla_cache() applied to the BF16-rounded roped rows.

Parameters:

  • kv (torch.Tensor) – Contiguous CUDA BF16 [num_tokens, 512] latent KV rows (unrotated).

  • cache (torch.Tensor) – Destination opaque uint8 NVFP4 paged cache (3-D, HND or NHD; any positive page size; 16-byte-multiple page stride). Only the addressed data rows and scale slots are written.

  • slot_mapping (torch.Tensor) – Contiguous 1-D CUDA int32 or int64 tensor with at most num_tokens entries; negative and out-of-range slots are skipped, and the caller may pass -1 for non-boundary rows instead of relying on compress_ratio. A valid slot must not repeat inside one call.

  • positions (torch.Tensor) – Contiguous 1-D CUDA int64 [num_tokens] positions; every inserted row’s cos/sin row must index inside cos_sin_cache (not checked on the device). int32 is rejected rather than converted.

  • cos_sin_cache (torch.Tensor) – Contiguous CUDA fp32 [max_position, 64] table, cos[32] | sin[32].

  • compress_ratio (int) – 1 or 2: the boundary rule and the cos/sin row rule above.

Returns:

The cache is updated in place; num_tokens == 0 or an empty slot_mapping returns without a launch. CUDA-Graph safe (no allocation, no host synchronisation).

Return type:

None