The local variable batch_size_ is initialized with query.size(0), which corresponds to the total number of query tokens (sum_seq_q), not the batch size. This variable shadows the batch_size function parameter, which holds the correct number of sequences. This is a critical bug.

Please remove this line and use the batch_size function parameter in the call to trtllm_paged_attention_context_launcher on line 260.

The batch_size_ variable used here is incorrect as it holds the total number of query tokens instead of the batch size. Please use the batch_size function parameter instead.

After this change, please also review the initialization of runner_params.mMaxSeqLenQ inside trtllm_paged_attention_context_launcher. It is currently set to batch_size, which would become the number of sequences after the fix above. This is likely incorrect, as it should be the maximum query sequence length in the batch.

        batch_size, max_seq_len, num_qo_heads, num_kv_heads, head_dim, page_size,

The parameter runner_params.mNumPagesInMemPool is hardcoded to 0. The previous implementation calculated this value based on the total device memory. Setting it to 0 could cause issues if the underlying kernel relies on this value for resource allocation or performance tuning.

Since key_value_cache is now a raw pointer, its size is not available here. Consider passing the total number of pages in the cache pool as an argument from the Python side to correctly calculate this value. For example, in csrc/trtllm_mla_kernel_launcher.cu, this was calculated as key_value_cache.size(0) * 2.

Similar to the decode launcher, runner_params.mNumPagesInMemPool is hardcoded to 0. This could be problematic. Please consider passing the total number of pages from the caller to compute a more meaningful value for this parameter.

The function signature of trtllm_paged_attention_mla_launcher still uses at::Tensor arguments, while trtllm_paged_attention_decode_launcher and trtllm_paged_attention_context_launcher in trtllm_fmha_kernel_launcher.cu have been refactored to use raw pointers.

For consistency and to decouple the CUDA kernels from PyTorch's at::Tensor, consider applying the same refactoring to this function as well.

Adding a default constructor to memset the struct is a good improvement for ensuring all members are zero-initialized. This follows RAII principles and makes the code safer and cleaner.

If there are other restrictions on group size (i.e., num_qo_heads // num_kv_heads) we should throw that error as well

it is added by one because flashinfer assumes that the sliding window attention should consider the extra token during masking, right ? probably we can add a comment here.

is the rounding for 16B alignment ? then what about the workspace_buffer

@PerkzZheng workspace_buffer is allocated by users. We should ask users to make sure workspace_buffer is 16B aligned. I have added comments for this.

thanks.

nit (P1): add documentation:

//! \brief Helper function to launch a trtllm paged attention kernel.
//! \note This function should not be called directly from another file. Use `trtllm_paged_attention_decode` and `trtllm_paged_attention_context` instead.
//!
//! \param out Device pointer to the output tensor.
//! \param query Device pointer to the input query tensor.
//! \param key_cache Device pointer to the input paged key cache tensor. The strides can be set with \p kv_stride_0 \p kv_stride_1 and \p kv_stride_2 .
//! \param value_cache Device pointer to the input paged value cache tensor. The strides can be set with \p kv_stride_0 \p kv_stride_1 and \p kv_stride_2 .
//! \param workspace_buffer Device pointer to the workspace. Must be at least 16-byte aligned. Recommended to allocate at least 128MB for workspace.
//! \param block_tables Device pointer to the block tables. The table shape is [batch_size, max_num_blocks_per_seq].
//! \param seq_lens Device pointer to the sequeunce lengths. The shape is [batch_size].
//! \param batch_size Batch size, i.e. the number of sequences in the batch.
//! \param max_q_len Maximum number of query tokens per sequence in the batch.
//! \param max_kv_len Maximum number of key/value tokens per sequence in the batch.
//! \param num_pages Maximum number of pages of the kv-cache.
//! \param num_qo_heads Number of query heads.
//! \param num_kv_heads Number of key/value heads.
//! \param head_dim_qk Head dimension of query/key.
//! \param head_dim_vo Head dimension of value/output.
//! \param page_size Number of tokens per page.
//! \param kv_stride_0 Stride of the "page_size" dimension of kv-cache with shape [num_pages, 2, num_kv_heads, page_size, head_dim].
//! \param kv_stride_1 Stride of the "num_kv_heads" dimension of kv-cache with shape [num_pages, 2, num_kv_heads, page_size, head_dim].
//! \param kv_stride_2 Stride of the "2" dimension of kv-cache with shape [num_pages, 2, num_kv_heads, page_size, head_dim].
//! \param max_num_blocks_per_seq Maximum number of blocks that can be allocated for a sequence.
//! \param bmm1_scale The scaling factor applied between BMM1 and Softmax, not including the LOG2E factor.
//! \param bmm2_scale The scaling factor applied after BMM2, not including the scaling factor for Softmax output.
//! \param window_left The window size for sliding window attention. Set to -1 to disable sliding window attention.
//! \param sum_seq_q Total number of query tokens within the batch.
//! \param sm_count The number of SMs on the GPU.
//! \param stream The cuda stream to launch the kernel on.
//! \param cum_seq_lens_q Device pointer to the tensor storing the accumulated sequence lengths of the query tokens in the batch. Not used in ForGen mode.
//! \param cum_seq_lens_kv Device pointer to the tensor storing the accumulated sequence lengths of the key/value tokens in the batch. Not used in ForGen mode.

Since std::unordered_map never invalidates references and cache outlives everything, we can do

static std::unordered_map<Key, TllmGenFmhaRunner, KeyHash> cache;

and return TllmGenFmhaRunner& instead

      auto runner = TllmGenFmhaRunner(q_data_type, kv_data_type, o_data_type);
      auto [it, ok] = cache.emplace(key, std::move(runner));
      return it->second;