flashinfer.mla.cake_dsv4_nvfp4_rope_quantize_insert

flashinfer.mla.cake_dsv4_nvfp4_rope_quantize_insert(q: Tensor, kv: Tensor, cache: Tensor, slot_mapping: Tensor, positions: Tensor, cos_sin_cache: Tensor, *, q_head_padded: int = 0, apply_q_rope: bool = True, q_inplace: bool = False) → Tensor

Rotate q, then rotate, quantize and insert kv into the DeepSeek-V4 NVFP4 cache.

One launch replaces the torch-level GPT-J RoPE of the query and latent KV, the head padding of the query and flashinfer.mla.nvfp4_quantize_append_sparse_mla_cache(): the cache bytes are identical to appending the BF16-rounded roped row, and q_out is the fp32 rotation of each live head rounded to BF16 with zero-filled padded heads. Rows >= slot_mapping.numel() (data-parallel padding) get their q_out rows but are not inserted; negative and out-of-range slots are skipped.

Parameters:

  • q (torch.Tensor) – Contiguous CUDA BF16 [num_tokens, num_heads, 512] queries (448 NoPE dims, then 64 RoPE dims), head-major.

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

  • cache (torch.Tensor) – Destination opaque uint8 NVFP4 paged cache: 3-D [num_pages, page_size, 384], HND [num_pages, 1, page_size, 384] or NHD [num_pages, page_size, 1, 384], any positive page size, 16-byte-multiple page stride (padded pools accepted). 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; slot_mapping[i] = page_id * page_size + entry_id for row i. Negative and out-of-range entries are padding and skipped; a valid slot must not repeat inside one call.

  • positions (torch.Tensor) – Contiguous 1-D CUDA int64 [num_tokens] positions indexing cos_sin_cache (positions[i] < cos_sin_cache.shape[0] is a caller precondition; it is not checked on the device). int32 is rejected rather than converted, so the hot path stays free of hidden launches.

  • cos_sin_cache (torch.Tensor) – Contiguous CUDA fp32 [max_position, 64] table, cos[32] | sin[32] per position (pair p of the 64 RoPE dims uses column p).

  • q_head_padded (int) – Padded head count of q_out: one of 8, 16, 32, 64, 128 and at least num_heads; 0 (default) writes the cache only – identical to cake_dsv4_nvfp4_kv_rope_quantize_insert() with compress_ratio=1 – and returns an empty [num_tokens, 0, 512] tensor.

  • apply_q_rope (bool) – Rotate the live query heads (default). False copies them unchanged (padded heads are still zero-filled).

  • q_inplace (bool) – Rotate the 64 RoPE dims of q in place and return q itself (no padded copy; the NoPE dims are neither read nor written). Requires q_head_padded to be 0 or num_heads (no head padding). This is the fast path for serving geometries whose local head count is already a supported padded count (every power-of-two tensor-parallel split of DeepSeek-V4). With apply_q_rope=False only the cache is written and q is returned untouched.

Returns:

q_out – a new BF16 [num_tokens, q_head_padded, 512] tensor (torch.empty, fully written), or q itself with q_inplace. The cache is updated in place. num_tokens == 0 returns the empty tensor (or q) without a launch.

Return type:

torch.Tensor