[wip] xqa nhd/hnd

upd

Author: yzh119

csrc/flashinfer_xqa_binding.cu

@@ -18,14 +18,10 @@
 
 #if MLA_WRAPPER
 void xqa_wrapper_mla(int64_t multiProcessorCount, double qScale, TensorView output, TensorView q,
-#if PAGED_KV_CACHE_LAYOUT == 1
-                     TensorView kCacheVLLM, TensorView vCacheVLLM,
-#else
-                     TensorView pool,
-#endif
-                     TensorView kvCachePageList, int64_t maxSeqLen, TensorView seqLen,
-                     int64_t batchSize, TensorView kvCacheScale, TensorView semaphores,
-                     TensorView scratch, bool enable_pdl);
+                     TensorView kCacheVLLM, TensorView vCacheVLLM, TensorView kvCachePageList,
+                     int64_t maxSeqLen, TensorView seqLen, int64_t batchSize,
+                     TensorView kvCacheScale, TensorView semaphores, TensorView scratch,
+                     bool enable_pdl);
 
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(xqa_wrapper_mla, xqa_wrapper_mla);
 
@@ -36,14 +32,9 @@ void xqa_wrapper(bool run_sm90_fp8_mha, int64_t multiProcessorCount, int64_t nbK
 #if LOW_PREC_OUTPUT
                  TensorView rcpOutScale,
 #endif
-                 TensorView q, tvm::ffi::Optional<TensorView> attentionSinks,
-#if PAGED_KV_CACHE_LAYOUT == 1
-                 TensorView kCacheVLLM, TensorView vCacheVLLM,
-#else
-                 TensorView pool,
-#endif
-                 TensorView kvCachePageList, int64_t maxSeqLen, TensorView seqLen,
-                 int64_t batchSize, TensorView kvCacheScale,
+                 TensorView q, tvm::ffi::Optional<TensorView> attentionSinks, TensorView kCacheVLLM,
+                 TensorView vCacheVLLM, TensorView kvCachePageList, int64_t maxSeqLen,
+                 TensorView seqLen, int64_t batchSize, TensorView kvCacheScale,
 #if SPEC_DEC
                  int64_t qSeqLen, TensorView qCuSeqLens, TensorView mask,
 #endif

csrc/xqa/defines.h

@@ -92,21 +92,6 @@ static_assert(SPEC_DEC, "SPEC_Q_SEQ_LEN should only be used when SPEC_DEC is ena
 #define TOKENS_PER_PAGE 32
 #endif
 
-// don't modify
-#ifndef USE_PAGED_KV_CACHE
-#define USE_PAGED_KV_CACHE (TOKENS_PER_PAGE > 0)
-#endif
-
-// Paged KV Cache Format
-// 0 - XQA Original
-// 1 - separate K and V cache pools, each with layout (batch, seq_len, head, head_elem) for
-// VLLM/SGLang
-#ifdef USE_PAGED_KV_CACHE
-#ifndef PAGED_KV_CACHE_LAYOUT
-#define PAGED_KV_CACHE_LAYOUT 0
-#endif
-#endif
-
 // don't modify
 #define USE_BEAM_SEARCH (BEAM_WIDTH > 1)
 
@@ -170,8 +155,7 @@ static_assert(CACHE_ELEM_ENUM != 0);
 #endif
 
 // true should be better if warpTile.x * cacheElemSize < 128. otherwise use false.
-#define GRP_LOAD_V \
-  (CACHE_ELEM_ENUM != 0) || (HEAD_ELEMS == 256 && USE_PAGED_KV_CACHE && BEAM_WIDTH > 1)
+#define GRP_LOAD_V (CACHE_ELEM_ENUM != 0) || (HEAD_ELEMS == 256 && BEAM_WIDTH > 1)
 
 // use custom barrier for NVRTC to avoid pulling in many headers
 #ifndef USE_CUSTOM_BARRIER