Enable stride aware indexing on top of strided blocks for block_size and heads (assumes contiguous head dim) of fused_qk_norm_rope_cache_pts_quant_shuffle()

[jhu960213]

Motivation

Enabled stride-aware indexing for the token and head dimensions of the KV cache in fused_qk_norm_rope_cache_pts_quant_shuffle, to support the newly altered KV cache layout in vLLM where K/V are packed along the head dim as [num_blocks, num_heads_kv, block_size, 2*head_size] (K = [..., :head_size], V = [..., head_size:]). The existing blocks-first ([num_blocks, 2, …]) and KV-first ([2, num_blocks, …]) layouts remain supported, this adds a 3rd layout on top of them. The kernel now reads the cache's real block/token/head strides instead of assuming a contiguous per-block arrangement; only the innermost head_size dim must stay contiguous. Zero strides fall back to the previous contiguous formulas, so existing callers are unchanged.

Test Plan

python3 op_tests/test_fused_qk_norm_rope_cache_quant.py -> should all pass.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4312 --add-label <label>