feat: add HND KV layout support

benchmark/benchmark_cutlass_flash_attn_decode.py

@@ -4,6 +4,7 @@
 
 # isort: off
 import gc
+import math
 
 import torch
 import triton
@@ -49,7 +50,7 @@ def calculate_memory_usage(q_len_sum, kv_len_sum, num_heads, head_size,
     return (query_memory + kv_cache_memory + output_memory) / (1000**3)
 
 
-def make_decode_with_paged_kv_input(config):
+def make_decode_with_paged_kv_input(config, kv_layout="NHD"):
     seq_lens, num_heads, head_size, block_size, \
     output_dtype, _, num_blocks, _, q_dtype, is_sink = config
     # if num_heads == (16, 1) and head_size == 256:
@@ -68,12 +69,26 @@ def make_decode_with_paged_kv_input(config):
                         num_query_heads,
                         head_size,
                         dtype=output_dtype)
-    key_cache = torch.randn(num_blocks,
-                            block_size,
-                            num_kv_heads,
-                            head_size,
-                            dtype=output_dtype)
-    value_cache = torch.randn_like(key_cache)
+    # Create KV cache in NHD layout first, then permute to HND if requested.
+    # HND: [num_blocks, num_kv_heads, block_size, head_size] (non-contiguous
+    # strides that the kernel auto-detects via is_paged_kv_hnd_layout()).
+    # NHD: [num_blocks, block_size, num_kv_heads, head_size] (contiguous).
+    key_cache_nhd = torch.randn(num_blocks,
+                                block_size,
+                                num_kv_heads,
+                                head_size,
+                                dtype=output_dtype)
+    value_cache_nhd = torch.randn(num_blocks,
+                                  block_size,
+                                  num_kv_heads,
+                                  head_size,
+                                  dtype=output_dtype)
+    if kv_layout == "HND":
+        key_cache = key_cache_nhd.permute(0, 2, 1, 3)
+        value_cache = value_cache_nhd.permute(0, 2, 1, 3)
+    else:
+        key_cache = key_cache_nhd
+        value_cache = value_cache_nhd
     cu_query_lens = torch.tensor([0] + query_lens,
                                  dtype=torch.int32).cumsum(dim=0,
                                                            dtype=torch.int32)
@@ -384,7 +399,7 @@ def _mk_cfg(seq_lens, num_heads, block_size, name, head_size=128,
 ]
 
 
-def benchmark_batch_decode(config, iterations=200):
+def benchmark_batch_decode(config, iterations=200, kv_layout="NHD"):
     """Benchmark a single batch decode config with GPU-event timing."""
     (seq_lens, num_heads, head_size, block_size, dtype, soft_cap,
      num_blocks, fa_versions, q_dtype, is_sink, name) = config
@@ -394,7 +409,8 @@ def benchmark_batch_decode(config, iterations=200):
     (maybe_quantized_query, maybe_quantized_key_cache,
      maybe_quantized_value_cache, max_query_len, cu_query_lens,
      max_kv_len, seq_k, scale, block_tables, sink, _,
-     _, _, _, _) = make_decode_with_paged_kv_input(full_config)
+     _, _, _, _) = make_decode_with_paged_kv_input(full_config,
+                                                   kv_layout=kv_layout)
 
     num_seqs = int(seq_lens.split(",")[0])
     max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
@@ -493,3 +509,57 @@ def _run(i):
         clear_xpu_cache()
 
     print("=" * 80)
+
+    # ================================================================
+    # NHD vs HND KV Layout Comparison
+    # ================================================================
+    print("\n" + "=" * 115)
+    print("Batch Decode: NHD vs HND KV Layout Comparison")
+    print("=" * 115)
+    hdr = (f"{'config':<40} | {'batch':>5} {'kv_sum':>7} | "
+           f"{'NHD(us)':>9} {'NHD_BW':>8} | "
+           f"{'HND(us)':>9} {'HND_BW':>8} | "
+           f"{'ratio':>6} {'winner':>6}")
+    print(hdr)
+    print("-" * 115)
+
+    for cfg in BATCH_DECODE_CONFIGS:
+        name = cfg[-1]
+        seq_lens = cfg[0]
+        num_seqs = int(seq_lens.split(",")[0])
+        kv_lens = list(map(int, seq_lens.split(",")[2].split("+")))
+        kv_sum = sum(kv_lens)
+
+        nhd_us = nhd_bw = hnd_us = hnd_bw = float("nan")
+        try:
+            nhd_us, nhd_bw = benchmark_batch_decode(cfg, iterations=200,
+                                                     kv_layout="NHD")
+        except Exception as e:
+            print(f"{name:<40} | {num_seqs:>5} {kv_sum:>7} | "
+                  f"NHD ERROR: {str(e)[:30]}")
+        clear_xpu_cache()
+
+        try:
+            hnd_us, hnd_bw = benchmark_batch_decode(cfg, iterations=200,
+                                                     kv_layout="HND")
+        except Exception as e:
+            print(f"{name:<40} | {num_seqs:>5} {kv_sum:>7} | "
+                  f"HND ERROR: {str(e)[:30]}")
+        clear_xpu_cache()
+
+        if (not math.isnan(nhd_us) and not math.isnan(hnd_us)
+                and nhd_us > 0 and hnd_us > 0):
+            ratio = nhd_us / hnd_us
+            winner = "HND" if ratio > 1.01 else (
+                "NHD" if ratio < 0.99 else "~same")
+            print(f"{name:<40} | {num_seqs:>5} {kv_sum:>7} | "
+                  f"{nhd_us:>9.1f} {nhd_bw:>8.1f} | "
+                  f"{hnd_us:>9.1f} {hnd_bw:>8.1f} | "
+                  f"{ratio:>6.3f} {winner:>6}")
+        else:
+            print(f"{name:<40} | {num_seqs:>5} {kv_sum:>7} | "
+                  f"{'N/A':>9} {'N/A':>8} | "
+                  f"{'N/A':>9} {'N/A':>8} | "
+                  f"{'N/A':>6} {'N/A':>6}")
+
+    print("=" * 115)

benchmark/benchmark_cutlass_flash_attn_varlen.py

@@ -4,6 +4,7 @@
 
 # isort: off
 import gc
+import math
 
 import torch
 import triton
@@ -43,7 +44,7 @@ def calculate_flops(num_query_heads, query_lens, kv_lens, head_size,
     return total
 
 
-def make_varlen_with_paged_kv_input(config):
+def make_varlen_with_paged_kv_input(config, kv_layout="NHD"):
     num_seqs, query_lens, kv_lens, num_heads, head_size, \
         block_size, window_size, output_dtype, _, num_blocks, \
         _, q_dtype, is_sink, is_causal, is_paged, kv_dtype = config
@@ -63,17 +64,32 @@ def make_varlen_with_paged_kv_input(config):
                         head_size,
                         dtype=output_dtype)
     if is_paged:
-        key_cache = torch.randn(num_blocks,
-                                block_size,
-                                num_kv_heads,
-                                head_size,
-                                dtype=output_dtype)
+        # Create KV cache in NHD layout first, then permute to HND if requested.
+        # HND: [num_blocks, num_kv_heads, block_size, head_size]
+        #   (non-contiguous strides auto-detected by the kernel).
+        # NHD: [num_blocks, block_size, num_kv_heads, head_size] (contiguous).
+        key_cache_nhd = torch.randn(num_blocks,
+                                    block_size,
+                                    num_kv_heads,
+                                    head_size,
+                                    dtype=output_dtype)
+        value_cache_nhd = torch.randn(num_blocks,
+                                      block_size,
+                                      num_kv_heads,
+                                      head_size,
+                                      dtype=output_dtype)
+        if kv_layout == "HND":
+            key_cache = key_cache_nhd.permute(0, 2, 1, 3)
+            value_cache = value_cache_nhd.permute(0, 2, 1, 3)
+        else:
+            key_cache = key_cache_nhd
+            value_cache = value_cache_nhd
     else:
         key_cache = torch.randn(sum(kv_lens),
                                 num_query_heads,
                                 head_size,
                                 dtype=output_dtype)
-    value_cache = torch.randn_like(key_cache)
+        value_cache = torch.randn_like(key_cache)
 
     cu_query_lens = torch.tensor([0] + query_lens,
                                  dtype=torch.int32).cumsum(dim=0,
@@ -217,7 +233,8 @@ def benchmark_varlen_with_paged_kv(num_seqs,
                                    is_paged,
                                    kv_dtype,
                                    provider,
-                                   iterations=20):
+                                   iterations=20,
+                                   kv_layout="NHD"):
     maybe_quantized_query, maybe_quantized_key_cache, \
         maybe_quantized_value_cache, \
         max_query_len, cu_query_lens, max_kv_len, cu_kv_lens, \
@@ -228,15 +245,10 @@ def benchmark_varlen_with_paged_kv(num_seqs,
                 query_lens, kv_lens, num_heads, head_size,
                 block_size, window_size, output_dtype, soft_cap,
                 num_blocks, fa_versions, q_dtype, is_sink,
-                is_causal, is_paged, kv_dtype))
+                is_causal, is_paged, kv_dtype),
+                kv_layout=kv_layout)
     num_query_heads = num_heads[0]
 
-    print(f"Running config: {num_seqs, query_lens, kv_lens, \
-                              num_heads, head_size, block_size, \
-                              window_size, output_dtype, soft_cap, num_blocks, \
-                              fa_versions, q_dtype, is_sink, is_causal, \
-                              is_paged, kv_dtype}, Provider: {provider}",
-          flush=True)
     assert iterations > 5, \
     "Number of iterations should be greater than 5 to account for warmup"
 
@@ -520,3 +532,74 @@ def filter_configs(configs):
     save_path = ensure_save_path_exists(args.save_path)
     # Run performance benchmark
     benchmark.run(print_data=True, save_path=save_path)
+
+    # ================================================================
+    # Varlen Paged: NHD vs HND KV Layout Comparison
+    # ================================================================
+    paged_configs = [c for c in configs if c[14] is True]  # is_paged=True
+    if paged_configs:
+        print("\n" + "=" * 115)
+        print("Varlen Paged: NHD vs HND KV Layout Comparison")
+        print("=" * 115)
+        hdr = (f"{'config':<50} | {'seqs':>4} {'kv_sum':>7} | "
+               f"{'NHD(us)':>9} {'HND(us)':>9} | "
+               f"{'ratio':>6} {'winner':>6}")
+        print(hdr)
+        print("-" * 115)
+
+        for cfg in paged_configs:
+            (num_seqs, query_lens, kv_lens, num_heads, head_size,
+             block_size, window_size, output_dtype, soft_cap,
+             num_blocks, fa_versions, q_dtype, is_sink, is_causal,
+             is_paged, kv_dtype) = cfg
+
+            # Build a readable config name (similar to decode benchmark)
+            dtype_str = "bf16" if output_dtype == torch.bfloat16 else "fp16"
+            causal_str = "causal" if is_causal else "nocausal"
+            sink_str = "_sink" if is_sink else ""
+            name = (f"B{num_seqs}_{num_heads}_h{head_size}"
+                    f"_blk{block_size}_{dtype_str}_{causal_str}{sink_str}")
+
+            # Compute kv_sum from the kv_lens string
+            kv_lens_list = list(map(int, kv_lens.split(",")))
+            kv_sum = sum(kv_lens_list)
+
+            nhd_us = hnd_us = float("nan")
+            try:
+                nhd_us = benchmark_varlen_with_paged_kv(
+                    num_seqs, query_lens, kv_lens, num_heads, head_size,
+                    block_size, window_size, output_dtype, soft_cap,
+                    num_blocks, fa_versions, q_dtype, is_sink, is_causal,
+                    is_paged, kv_dtype, provider="flash_kernel_time",
+                    iterations=iterations, kv_layout="NHD")
+            except Exception as e:
+                print(f"{name:<50} | {num_seqs:>4} {kv_sum:>7} | "
+                      f"NHD ERROR: {str(e)[:30]}")
+            clear_xpu_cache()
+
+            try:
+                hnd_us = benchmark_varlen_with_paged_kv(
+                    num_seqs, query_lens, kv_lens, num_heads, head_size,
+                    block_size, window_size, output_dtype, soft_cap,
+                    num_blocks, fa_versions, q_dtype, is_sink, is_causal,
+                    is_paged, kv_dtype, provider="flash_kernel_time",
+                    iterations=iterations, kv_layout="HND")
+            except Exception as e:
+                print(f"{name:<50} | {num_seqs:>4} {kv_sum:>7} | "
+                      f"HND ERROR: {str(e)[:30]}")
+            clear_xpu_cache()
+
+            if (not math.isnan(nhd_us) and not math.isnan(hnd_us)
+                    and nhd_us > 0 and hnd_us > 0):
+                ratio = nhd_us / hnd_us
+                winner = "HND" if ratio > 1.01 else (
+                    "NHD" if ratio < 0.99 else "~same")
+                print(f"{name:<50} | {num_seqs:>4} {kv_sum:>7} | "
+                      f"{nhd_us:>9.1f} {hnd_us:>9.1f} | "
+                      f"{ratio:>6.3f} {winner:>6}")
+            else:
+                print(f"{name:<50} | {num_seqs:>4} {kv_sum:>7} | "
+                      f"{'N/A':>9} {'N/A':>9} | "
+                      f"{'N/A':>6} {'N/A':>6}")
+
+        print("=" * 115)

csrc/flash_attn/flash_api.cpp

@@ -279,8 +279,11 @@ std::vector<at::Tensor> mha_varlen_fwd(
     int num_tokens = batch_size;
     int num_heads_q = q.size(1);
     int head_dim = q.size(2);
-    int num_heads_kv = k.size(2);
-    int kv_block_size = k.size(1);
+    // Layout-aware extraction: HND (permuted) has size(1)=num_heads,
+    // size(2)=block_size
+    bool is_hnd_k = is_paged_kv_hnd_layout(k);
+    int num_heads_kv = is_hnd_k ? k.size(1) : k.size(2);
+    int kv_block_size = is_hnd_k ? k.size(2) : k.size(1);
 
     int num_kv_splits = 1;
     at::Tensor tmp_out = out;
@@ -334,8 +337,11 @@ std::vector<at::Tensor> mha_varlen_fwd(
     int batch_size = static_cast<int>(cu_seqlens_q.size(0)) - 1;
     int num_heads_q = q.size(1);
     int v_head_dim = v.size(-1);
-    int num_heads_kv = k.size(2);
-    int block_size = k.size(1);
+    // Layout-aware extraction: HND (permuted) has size(1)=num_heads,
+    // size(2)=block_size
+    bool is_hnd_k = is_paged_kv_hnd_layout(k);
+    int num_heads_kv = is_hnd_k ? k.size(1) : k.size(2);
+    int block_size = is_hnd_k ? k.size(2) : k.size(1);
     int head_size_qk = q.size(-1);
 
     // SLM (shared local memory) limits on Intel Xe restrict the paged decode

csrc/xpu/attn/paged_kv_utils.h

@@ -3,13 +3,40 @@
 #include <limits>
 #include <torch/all.h>
 
-// Normalize the physical page stride to sequence-position units. For a regular
-// contiguous layout [num_blocks, block_size, num_heads, head_size], this equals
-// block_size. For interleaved or cross-layer layouts, it includes the physical
-// gaps between logical KV blocks.
+// Detect whether a paged KV cache tensor uses HND layout.
+//
+// paged KV supports two layouts:
+//   NHD: [num_blocks, block_size, num_heads, head_size]  (contiguous)
+//   HND: [num_blocks, num_heads, block_size, head_size]  (permuted from NHD)
+//
+// HND is created by permuting an NHD tensor via .permute(0, 2, 1, 3).
+// After permutation the strides encode the layout:
+//   NHD contiguous:  stride(1) = num_heads*head_size  > stride(2) = head_size
+//   HND (permuted):  stride(1) = head_size            < stride(2) =
+//   num_heads*head_size
+//
+// So the check simplifies to: stride(1) < stride(2) means dim1 and dim2
+// were swapped, i.e. HND layout.
+//
+// This check is robust for cross-layer (non-contiguous page stride) tensors
+// because stride(1) and stride(2) within each block are unchanged by the
+// outer non-contiguous stride(0).
+inline bool is_paged_kv_hnd_layout(const at::Tensor& key_cache) {
+  return key_cache.stride(2) > key_cache.stride(1);
+}
+
+// Normalize the physical page stride to sequence-position units.
+// paged KV: NHD [num_blocks, block_size, num_heads, head_size]
+//        or HND [num_blocks, num_heads, block_size, head_size] (permuted).
+// For a regular contiguous NHD layout this equals block_size.
+// For interleaved or cross-layer layouts it includes the physical gaps between
+// logical KV blocks.
 inline int64_t
 get_paged_kv_cache_page_stride_elements(const at::Tensor& key_cache) {
-  return key_cache.stride(0) / key_cache.stride(1);
+  // seq stride: NHD→stride(1), HND (permuted)→stride(2)
+  int64_t seq_stride = is_paged_kv_hnd_layout(key_cache) ? key_cache.stride(2)
+                                                         : key_cache.stride(1);
+  return key_cache.stride(0) / seq_stride;
 }
 
 // Return the sequence length needed by the 2D load surface to cover the full
@@ -23,11 +50,27 @@ get_paged_kv_cache_effective_total_seqlen(const at::Tensor& key_cache) {
 
 // Validate that the physical page stride can be converted to sequence-position
 // units and safely passed through int32 kernel parameters.
+// Supports both NHD [num_blocks, block_size, num_heads, head_size] and
+// HND [num_blocks, num_heads, block_size, head_size] (permuted) layouts.
 inline void check_paged_kv_cache_strides(
     const at::Tensor& key_cache, const at::Tensor& value_cache) {
-  int64_t k_stride_seq = key_cache.stride(1);
+  // Ensure K and V use the same layout — downstream kernels derive layout
+  // from key_cache only and apply it to both tensors.
+  bool k_is_hnd = is_paged_kv_hnd_layout(key_cache);
+  bool v_is_hnd = is_paged_kv_hnd_layout(value_cache);
+  TORCH_CHECK(
+      k_is_hnd == v_is_hnd,
+      "Paged K and V caches must use the same layout (both NHD or both HND), "
+      "but got K=",
+      k_is_hnd ? "HND" : "NHD",
+      " V=",
+      v_is_hnd ? "HND" : "NHD");
+
+  // seq stride: NHD→stride(1), HND (permuted)→stride(2)
+  int64_t k_stride_seq = k_is_hnd ? key_cache.stride(2) : key_cache.stride(1);
   int64_t k_physical_page_stride = key_cache.stride(0);
-  int64_t v_stride_seq = value_cache.stride(1);
+  int64_t v_stride_seq =
+      v_is_hnd ? value_cache.stride(2) : value_cache.stride(1);
   int64_t v_physical_page_stride = value_cache.stride(0);
   TORCH_CHECK(
       k_stride_seq > 0,

csrc/xpu/attn/xe_2/collective/chunk_prefill_mainloop.hpp

@@ -1001,7 +1001,7 @@ struct DecodeFwdMainloop<
                           next_tile_idx % tiles_per_page;
         } else {
           // set to last page
-          next_tile_idx = params.max_pages_per_seq * tiles_per_page - 1;
+          next_tile_idx = params.max_pages_per_seq * page_stride_tiles - 1;
         }
       }
       tile_idx = next_tile_idx;