docs(base): add vLLM interface references to FlashAttention, ReshapeAndCache, RotaryEmbedding

Author: zhangyue

src/base/flash_attention.h

@@ -9,6 +9,14 @@
 
 namespace infini::ops {
 
+// Fused multi-head / grouped-query attention.
+//
+// Interface follows vLLM v1 `AttentionImpl.forward()`:
+//   `vllm.v1.attention.backends.abstract.AttentionImpl`
+//
+// Layout: `query` / `key` / `value` are `[T, N, D]` (TND).
+// Prefill uses `cu_seqlens_q` / `cu_seqlens_kv` for variable-length packing.
+// Decode uses `block_table` for paged KV cache lookup.
 class FlashAttention : public Operator<FlashAttention> {
  public:
   FlashAttention(const Tensor query, const Tensor key, const Tensor value,

src/base/reshape_and_cache.h

@@ -8,6 +8,15 @@
 
 namespace infini::ops {
 
+// Scatter `key` / `value` tokens into a paged KV cache.
+//
+// Interface follows vLLM's `reshape_and_cache` kernel:
+//   `vllm._custom_ops.reshape_and_cache_flash`
+//
+// `kv_cache` layout: `[2, num_blocks, block_size, num_kv_heads, head_size]`.
+// `slot_mapping`: 1D `[num_tokens]`, each entry is the linear slot index
+// into the cache.  Padding tokens must be filtered by the caller (no
+// negative indices).
 class ReshapeAndCache : public Operator<ReshapeAndCache> {
  public:
   ReshapeAndCache(const Tensor key, const Tensor value, const Tensor kv_cache,

src/base/rotary_embedding.h

@@ -8,6 +8,15 @@
 
 namespace infini::ops {
 
+// Rotary position embedding (RoPE) applied in-place to Q and K.
+//
+// Interface follows vLLM's `RotaryEmbedding.forward_oot()`:
+//   `vllm.model_executor.layers.rotary_embedding.RotaryEmbedding`
+//
+// `positions`: `[T]` token position indices.
+// `cos_sin_cache`: precomputed `[max_seq_len, rotary_dim]` table.
+// `query` / `key`: `[T, N, D]` (TND layout), mutated in-place into
+// `query_out` / `key_out`.
 class RotaryEmbedding : public Operator<RotaryEmbedding> {
  public:
   RotaryEmbedding(const Tensor positions, const Tensor query, const Tensor key,