flashinfer.concat_ops.concat_mla_k

flashinfer.concat_ops.concat_mla_k(k: Tensor, k_nope: Tensor, k_rope: Tensor) → None

Concatenate k_nope and k_rope tensors for MLA attention.

This function efficiently concatenates:

  • k_nope: per-head nope values

  • k_rope: shared rope values (broadcast to all heads)

Supported dtypes: torch.bfloat16, torch.float16,

torch.float8_e4m3fn, torch.float8_e5m2.

Key optimizations:

  • Warp-based processing with software pipelining

  • Vectorized memory access (compile-time dispatch per dtype)

  • L2 prefetching for next row while processing current

  • Register reuse for rope values across all heads in a chunk

Parameters:

  • k (torch.Tensor) – Output tensor, shape: [num_tokens, num_heads, nope_dim + rope_dim]. Modified in-place.

  • k_nope (torch.Tensor) – The nope part of k, shape: [num_tokens, num_heads, nope_dim].

  • k_rope (torch.Tensor) – The rope part of k (shared), shape: [num_tokens, 1, rope_dim]. This is broadcast to all heads.

Example

>>> import torch
>>> import flashinfer
>>> num_tokens = 2048
>>> num_heads = 128
>>> nope_dim = 128
>>> rope_dim = 64
>>> # BF16 example
>>> k = torch.empty(num_tokens, num_heads, nope_dim + rope_dim, dtype=torch.bfloat16, device="cuda")
>>> k_nope = torch.randn(num_tokens, num_heads, nope_dim, dtype=torch.bfloat16, device="cuda")
>>> k_rope = torch.randn(num_tokens, 1, rope_dim, dtype=torch.bfloat16, device="cuda")
>>> flashinfer.concat_ops.concat_mla_k(k, k_nope, k_rope)

Note

This kernel is specifically optimized for: - num_heads = 128 - nope_dim = 128 - rope_dim = 64