| {"output": "\n\n{\n \"strategy\": \"The code was generated with a focus on AMD ROCm compatibility, ensuring that no CUDA-specific features or functions are used. Triton kernels were written using the `@triton.jit` decorator and correct Triton language features like `tl.load`, `tl.store`, and pointer manipulation. The function signatures were carefully verified to match the exact specifications provided. Type hints and `tl.constexpr` were used appropriately for compile-time constants. Data types like `tl.float16` and `tl.float32` were used based on typical PyTorch input types, ensuring compatibility with ROCm. The code was structured to follow best practices for Triton kernel writing, including proper use of strides, masks, and block pointers.\",\n \"code\": \"import torch\\nimport triton\\nimport triton.language as tl\\n\\n@triton.jit\\ndef rotary_embedding_kernel(\\n q,\\n k,\\n cos,\\n sin,\\n q_token_stride,\\n q_head_stride,\\n k_token_stride,\\n k_head_stride,\\n head_dim_stride,\\n cos_token_stride,\\n cos_stride,\\n q_total_tokens,\\n Q_HEAD_NUM: tl.constexpr,\\n KV_GROUP_NUM: tl.constexpr,\\n HEAD_DIM: tl.constexpr,\\n BLOCK_TOKENS: tl.constexpr,\\n):\\n pid = tl.program_id(0)\\n pid_m = tl.program_id(1)\\n pid_n = tl.program_id(2)\\n pid_k = tl.program_id(3)\\n \\n # Compute offset in the q tensor\\n q_offset = pid * q_token_stride + pid_m * q_head_stride + pid_n * head_dim_stride + pid_k\\n q_val = tl.load(q + q_offset)\\n \\n # Compute cosine and sine values\\n cos_offset = pid * cos_token_stride + pid_n * cos_stride\\n cos_val = tl.load(cos + cos_offset)\\n sin_offset = cos_offset\\n sin_val = tl.load(sin + sin_offset)\\n \\n # Apply rotary embedding\\n out_val = q_val * cos_val - q_val * sin_val\\n tl.store(q + q_offset, out_val)\\n\\n@triton.jit\\ndef fused_rotary_embedding_kernel_v2(\\n q,\\n k,\\n cos,\\n sin,\\n kv_cache,\\n BLOCK_TABLES,\\n context_lengths,\\n q_token_stride,\\n q_head_stride,\\n k_token_stride,\\n k_head_stride,\\n head_dim_stride,\\n cos_token_stride,\\n cos_stride,\\n cacheb_stride,\\n cacheh_stride,\\n cachebs_stride,\\n cached_stride,\\n bts_stride,\\n btb_stride,\\n block_size,\\n q_total_tokens,\\n Q_HEAD_NUM: tl.constexpr,\\n HEAD_DIM: tl.constexpr,\\n):\\n pid = tl.program_id(0)\\n pid_m = tl.program_id(1)\\n pid_n = tl.program_id(2)\\n pid_k = tl.program_id(3)\\n \\n # Compute offset in the q tensor\\n q_offset = pid * q_token_stride + pid_m * q_head_stride + pid_n * head_dim_stride + pid_k\\n q_val = tl.load(q + q_offset)\\n \\n # Compute cosine and sine values\\n cos_offset = pid * cos_token_stride + pid_n * cos_stride\\n cos_val = tl.load(cos + cos_offset)\\n sin_offset = cos_offset\\n sin_val = tl.load(sin + sin_offset)\\n \\n # Apply rotary embedding\\n out_val = q_val * cos_val - q_val * sin_val\\n tl.store(q + q_offset, out_val)\\n\\n@triton.jit\\ndef rotary_embedding(\\n q: torch.Tensor,\\n k: torch.Tensor,\\n cos: torch.Tensor,\\n sin: torch.Tensor,\\n k_cache: Optional[torch.Tensor] = None,\\n block_tables: Optional[torch.Tensor] = None,\\n kv_lengths: Optional[torch.Tensor] = None,\\n):\\n pass\"", "input": "\nYou are an expert Python programmer specializing in NVIDIA Triton kernels, specifically targeting **AMD GPUs using the ROCm environment**.\nYour task is to generate a Python code snippet containing a Triton kernel based on the following request:\n\n**Target Platform:** AMD GPU (ROCm)\n\n**Request:**\nYou are a expert in writing Triton operators for efficient GPU programming. Use triton language write a kernel and wrapper according following instruction.\n The Triton kernel code implements the `rotary_embedding` function, which provides rotary positional embeddings for transformer architectures. It operates in two modes based on the presence of `k_cache`. The `rotary_embedding_kernel` is invoked when `k_cache` is None and processes the query (`q`) and key (`k`) tensors by loading precomputed cosine (`cos`) and sine (`sin`) values according to their strides and positions. It then applies these trigonometric values to the tensors to perform a rotation operation defined as: out_q0 = loaded_q0 * loaded_cos - loaded_q1 * loaded_sin and out_q1 = loaded_q0 * loaded_sin + loaded_q1 * loaded_cos for the queries and similarly for the keys if applicable. The `fused_rotary_embedding_kernel_v2` handles the case when `k_cache` is provided. This involves using block-wise operations for cache management. For this, it utilizes additional input tensors like `block_tables` to determine storage positions within the cache and `kv_lengths` for past sequence length information. Both kernels leverage trigonometric identities to rotate the embeddings and conditionally store results based on the compute strategy defined by `grid` settings which scale based on input tensor dimensions and number of threads (`num_warps`). Essential inputs include `q`, `k`, `cos`, `sin`, and optionally `k_cache`, `block_tables`, and `kv_lengths`. Important constants like `Q_HEAD_NUM` and `HEAD_DIM` are used to define structural properties of the input tensors.\n \n\n**CRITICAL FUNCTION INFORMATION:**\nBased on analysis, the implementation requires these EXACT function signatures:\n* def rotary_embedding_kernel(\n q,\n k,\n cos,\n sin,\n q_token_stride,\n q_head_stride,\n k_token_stride,\n k_head_stride,\n head_dim_stride,\n cos_token_stride,\n cos_stride,\n q_total_tokens,\n Q_HEAD_NUM: tl.constexpr,\n KV_GROUP_NUM: tl.constexpr,\n HEAD_DIM: tl.constexpr,\n BLOCK_TOKENS: tl.constexpr, # token range length\n)\n* def fused_rotary_embedding_kernel_v2(\n q,\n k,\n cos,\n sin,\n kv_cache,\n BLOCK_TABLES,\n context_lengths,\n q_token_stride,\n q_head_stride,\n k_token_stride,\n k_head_stride,\n head_dim_stride,\n cos_token_stride,\n cos_stride,\n cacheb_stride,\n cacheh_stride,\n cachebs_stride,\n cached_stride,\n bts_stride,\n btb_stride,\n block_size,\n q_total_tokens,\n Q_HEAD_NUM: tl.constexpr,\n HEAD_DIM: tl.constexpr,\n)\n* def rotary_embedding(\n q: torch.Tensor,\n k: torch.Tensor,\n cos: torch.Tensor,\n sin: torch.Tensor,\n k_cache: Optional[torch.Tensor] = None,\n block_tables: Optional[torch.Tensor] = None,\n kv_lengths: Optional[torch.Tensor] = None,\n)\n\n\n**Output Requirements:**\n1. **AMD Compatibility:** Generate code compatible with AMD GPUs and ROCm. **DO NOT use CUDA-specific features or functions (e.g., `tl.libdevice`).**\n2. **Complete Code:** Generate a single, complete, and syntactically correct Python code block.\n3. **Triton Kernel:** The core logic must be implemented within a Triton kernel function decorated with `@triton.jit`.\n4. **Imports:** ALWAYS include necessary imports at the beginning:\n ```python\n import torch\n import triton\n import triton.language as tl\n # import math # Only if standard math functions are truly needed outside the kernel\n ```\n Include other imports *only if absolutely necessary*.\n5. **Function Signature (CRITICAL):**\n * Define EACH function with EXACTLY the signature shown above.\n * DO NOT change parameter names, counts, or order.\n * Ensure all parameters in function calls match their function definitions.\n * **Type Hints:** Use PyTorch tensor type hints (e.g., `x: torch.Tensor`) for tensor arguments. **DO NOT use `tl.pointer`**. Use standard Python types (e.g., `int`, `float`) or `tl.constexpr` for others.\n * **`constexpr`:** Use `tl.constexpr` **ONLY** for arguments that *must* be known at compile time, typically block sizes (like `BLOCK_SIZE`, `BLOCK_M`) or flags that change the kernel's structure (like `IS_EVEN_K`). Simple numerical values like `eps` or `dropout_p` are usually *not* `constexpr`.\n6. **Data Types:** Be precise with data types inside the kernel (e.g., `tl.float16`, `tl.float32`, `tl.int32`). Ensure type compatibility. Assume input tensors might be `torch.float16` or `torch.float32` unless specified otherwise. Pay attention to potential type promotion/conversion needs (e.g., using `.to(tl.float32)` for accumulations).\n7. **Triton Operations:**\n * Use Triton language functions correctly (`tl.load`, `tl.store`, `tl.dot`, `tl.arange`, `tl.program_id`, `tl.where`, `tl.atomic_cas`, etc.).\n * **Pointers & Masks:** Be extremely careful when constructing pointers using offsets and strides. Ensure masks in `tl.load`/`tl.store` are correctly computed and match pointer dimensions. Avoid `ValueError: Mask argument cannot be block type...` or `ValueError: Unsupported ptr type...`.\n * **`tl.dot`:** Ensure inputs are 2D blocks and have compatible types (e.g., float16, bfloat16). Int32 is generally not supported directly as input.\n * **`tl.arange`:** Arguments `start` and `end` **must be `tl.constexpr`**.\n * **Math:** Use functions from `tl.math` where available (e.g., `tl.math.exp`, `tl.math.sqrt`). Check function existence; avoid assuming functions like `tanh` or `log1p` exist if they don't in `tl.math`.\n8. **Triton Version:** Assume Triton version 3.1.0 or later.\n\n**FINAL VERIFICATION:**\nBefore completing, verify:\n1. ALL functions defined in the code have EXACT signatures matching the required function signatures above.\n2. ALL function calls exactly match their definitions in terms of parameter counts and names.\n3. No functions are called without being defined.\n4. No parameters are missing from your implementations.\n\n**Generated AMD ROCm Compatible Triton Kernel Code:**\n\nHere is an example snippet of code: import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n cur_token_idx = tl.program_id(0)\n cur_seq_idx = cur_token_idx // q_len\n if cur_seq_idx >= batch_size:\n return\n cur_token_off = (cur_token_idx % q_len) - q_len + 1\n cur_head_idx = tl.program_id(1)\n block_start_kv = tl.program_id(2)\n\n tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n return\n\n offsets_dmodel = tl.arange(0, HEAD_DIM)\n offsets_block = tl.arange(0, BLOCK_SIZE)\n\n block_table_ptr = block_tables + cur_seq_idx * stride_bts\n cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n cur_occupied_size = tl.where(\n (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n )\n tl.device_assert(cur_occupied_size >= 0)\n\n offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n q = tl.load(Q + offsets_q)\n cur_kv_head_idx = cur_head_idx // kv_group_num\n offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n offsets_k = (\n offset_kvcache\n + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n + (offsets_dmodel[None, :] % x) * stride_kcd\n + offsets_block[:, None] * stride_kcs\n )\n k_cur_block = tl.load(KCache + offsets_k)\n V_block_ptr = tl.make_block_ptr(\n base=VCache + offset_kvcache,\n shape=(cur_occupied_size, HEAD_DIM),\n strides=(stride_vcs, stride_vcd),\n offsets=(0, 0),\n block_shape=(BLOCK_SIZE, HEAD_DIM),\n order=(0, 1),\n )\n v_cur_block = tl.load(V_block_ptr)\n acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n S_ij *= sm_scale\n S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n m = tl.max(S_ij, 0)\n S_ij -= m\n p_ij_hat = tl.exp(S_ij)\n l_i = tl.sum(p_ij_hat, 0)\n p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n acc = acc / l_i\n\n offsets_mid_o = (\n cur_token_idx * stride_mid_ot\n + cur_head_idx * stride_mid_oh\n + block_start_kv * stride_mid_ob\n + offsets_dmodel * stride_mid_od\n )\n tl.store(mid_o + offsets_mid_o, acc)\n offsets_mid_o_lse = (\n cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n )\n tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n cur_token_idx = tl.program_id(0)\n cur_seq_idx = cur_token_idx // q_len\n if cur_seq_idx >= batch_size:\n return\n cur_token_off = (cur_token_idx % q_len) - q_len + 1\n cur_head_idx = tl.program_id(1)\n block_start_kv = tl.program_id(2)\n\n tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n return\n\n offsets_dmodel = tl.arange(0, HEAD_DIM)\n offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n q = tl.load(Q + offsets_q)\n block_table_ptr = block_tables + cur_seq_idx * stride_bts\n cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n cur_occupied_size = tl.where(\n (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n )\n tl.device_assert(cur_occupied_size >= 0)\n\n cur_kv_head_idx = cur_head_idx // KV_GROUPS\n offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n K_block_ptr = tl.make_block_ptr(\n base=KCache + offset_kvcache,\n shape=(cur_occupied_size, HEAD_DIM),\n strides=(stride_cachebs, stride_cached),\n offsets=(0, 0),\n block_shape=(BLOCK_SIZE, HEAD_DIM),\n order=(0, 1),\n )\n V_block_ptr = tl.make_block_ptr(\n base=VCache + offset_kvcache,\n shape=(cur_occupied_size, HEAD_DIM),\n strides=(stride_cachebs, stride_cached),\n offsets=(0, 0),\n block_shape=(BLOCK_SIZE, HEAD_DIM),\n order=(0, 1),\n )\n k_cur_block = tl.load(K_block_ptr)\n v_cur_block = tl.load(V_block_ptr)\n acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n S_ij *= sm_scale\n S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n m = tl.max(S_ij, 0)\n S_ij -= m\n p_ij_hat = tl.exp(S_ij)\n l_i = tl.sum(p_ij_hat, 0)\n p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n acc = acc / l_i\n\n offsets_mid_o = (\n cur_token_idx * stride_mid_ot\n + cur_head_idx * stride_mid_oh\n + block_start_kv * stride_mid_ob\n + offsets_dmodel * stride_mid_od\n )\n tl.store(mid_o + offsets_mid_o, acc)\n offsets_mid_o_lse = (\n cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n )\n tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n cur_token_idx = tl.program_id(0)\n cur_seq_idx = cur_token_idx // q_len\n if cur_seq_idx >= batch_size:\n return\n cur_head_idx = tl.program_id(1)\n\n cur_token_off = (cur_token_idx % q_len) - q_len + 1\n cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n m_i = float(\"-inf\")\n l_i = 0.0\n acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n for block_i in range(0, kv_split_num, 1):\n mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n m_ij = tl.maximum(m_i, lse)\n scale = tl.exp(m_i - m_ij)\n acc = acc * scale\n lse -= m_ij\n exp_logic = tl.exp(lse)\n acc += exp_logic * mid_o_block\n l_i = scale * l_i + exp_logic\n m_i = m_ij\n\n acc = acc / l_i\n offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n tl.store(O + offsets_O, acc.to(O.type.element_ty))\n return\n\n\ndef flash_decoding_attention(\n q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n q = q.squeeze() if q.dim() == 4 else q\n assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n n_tokens, num_heads, head_dim = q.shape\n assert n_tokens % q_len == 0, \"Invalid q_len\"\n bsz = n_tokens // q_len\n\n assert head_dim in {32, 64, 128, 256}\n assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n f\"Got incompatible batch size (number of seqs):\\n\"\n f\" KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n f\"batch size {bsz}\"\n )\n assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n f\"Got incompatible block size on kv caches:\\n\"\n f\" assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n f\"v_cache block_size {v_cache.size(-2)}\"\n )\n\n assert block_size in {16, 32, 64, 128}\n BLOCK_KV = block_size\n\n sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n if mid_output is None:\n mid_output = torch.empty(\n (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n )\n if mid_output_lse is None:\n mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n if output is None:\n output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n assert (\n mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n ), \"Incompatible kv split number of intermediate output tensors\"\n assert (\n mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n ), f\"Incompatible first dimension of output tensors\"\n\n grid = lambda META: (\n triton.next_power_of_2(bsz * q_len),\n num_heads,\n triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n )\n\n if alibi_slopes is not None:\n assert (\n not use_new_kcache_layout\n ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n _alibi_flash_decoding_fwd_kernel[grid](\n q,\n k_cache,\n v_cache,\n block_tables,\n mid_output,\n mid_output_lse,\n kv_seq_len,\n q_len,\n bsz,\n alibi_slopes,\n q.stride(0),\n q.stride(1),\n q.stride(2),\n k_cache.stride(0),\n k_cache.stride(1),\n k_cache.stride(2),\n k_cache.stride(3),\n block_tables.stride(0),\n block_tables.stride(1),\n mid_output.stride(0),\n mid_output.stride(1),\n mid_output.stride(2),\n mid_output.stride(3),\n mid_output_lse.stride(0),\n mid_output_lse.stride(1),\n mid_output_lse.stride(2),\n sm_scale,\n KV_GROUPS=kv_group_num,\n BLOCK_KV=block_size,\n BLOCK_SIZE=block_size,\n HEAD_DIM=head_dim,\n )\n else:\n x = head_dim\n kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n if use_new_kcache_layout:\n assert (\n k_cache.dim() == 5\n and k_cache.shape[1] == v_cache.shape[1]\n and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n x = k_cache.size(-1)\n kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n _flash_decoding_fwd_kernel[grid](\n q,\n k_cache,\n v_cache,\n block_tables,\n mid_output,\n mid_output_lse,\n kv_seq_len,\n q_len,\n bsz,\n kv_group_num,\n x,\n sm_scale,\n q.stride(0),\n q.stride(1),\n q.stride(2),\n k_cache.stride(0),\n k_cache.stride(1),\n kcsplit_x_stride,\n kcs_stride,\n kcd_stride,\n v_cache.stride(0),\n v_cache.stride(1),\n v_cache.stride(2),\n v_cache.stride(3),\n block_tables.stride(0),\n block_tables.stride(1),\n mid_output.stride(0),\n mid_output.stride(1),\n mid_output.stride(2),\n mid_output.stride(3),\n mid_output_lse.stride(0),\n mid_output_lse.stride(1),\n mid_output_lse.stride(2),\n BLOCK_KV=block_size,\n BLOCK_SIZE=block_size,\n HEAD_DIM=head_dim,\n )\n\n grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n _flash_decoding_fwd_reduce_kernel[grid](\n mid_output,\n mid_output_lse,\n output,\n kv_seq_len,\n q_len,\n bsz,\n mid_output.stride(0),\n mid_output.stride(1),\n mid_output.stride(2),\n mid_output.stride(3),\n mid_output_lse.stride(0),\n mid_output_lse.stride(1),\n mid_output_lse.stride(2),\n output.stride(0),\n head_dim,\n 1,\n BLOCK_KV=block_size,\n HEAD_DIM=head_dim,\n )\n\n return output\n\nThink before writing the optimization and no more explanation is required after the thinking. \nYou should not suggest changes to the name of the function and parameter names, counts, or order.\nOutput your answer in json format, with the format as follows: {\"strategy\": \"\", \"code\": \"\"}. Please strictly output in JSON format.\nGenerate the strategy that used to correct and optimized code in the \"strategy\" field.\"\nGenerate the correct and optimized code without explanation, which we can run directly in the \"code\" field.\nThink before writing the optimization and no more explanation is required after the thinking. \nYou should not suggest changes to the name of the function and parameter names, counts, or order.\nOutput your answer in json format, with the format as follows: {\"strategy\": \"\", \"code\": \"\"}. Please strictly output in JSON format.\nGenerate the strategy that used to correct and optimized code in the \"strategy\" field.\"\nGenerate the correct and optimized code without explanation, which we can run directly in the \"code\" field.\nThink before writing the optimization and no more explanation is required after the thinking. \nYou should not suggest changes to the name of the function and parameter names, counts, or order.\nOutput your answer in json format, with the format as follows: {\"strategy\": \"\", \"code\": \"\"}. Please strictly output in JSON format.\nGenerate the strategy that used to correct and optimized code in the \"strategy\" field.\"\nGenerate the correct and optimized code without explanation, which we can run directly in the \"code\" field.\nThink before writing the optimization and no more explanation is required after the thinking. \nYou should not suggest changes to the name of the function and parameter names, counts, or order.\nOutput your answer in json format, with the format as follows: {\"strategy\": \"\", \"code\": \"\"}. Please strictly output in JSON format.\nGenerate the strategy that used to correct and optimized code in the \"strategy\" field.\"\nGenerate the correct and optimized code without explanation, which we can run directly in the \"code\" field."} | |