diff --git a/deep_gemm/include/deep_gemm/impls/sm120_fp8_mqa_logits.cuh b/deep_gemm/include/deep_gemm/impls/sm120_fp8_mqa_logits.cuh
index 93fc69daba..3987b081da 100644
--- a/deep_gemm/include/deep_gemm/impls/sm120_fp8_mqa_logits.cuh
+++ b/deep_gemm/include/deep_gemm/impls/sm120_fp8_mqa_logits.cuh
@@ -58,7 +58,8 @@ void sm120_fp8_mqa_logits(const uint32_t seq_len, const uint32_t seq_len_kv,
     DG_STATIC_ASSERT(BLOCK_KV == kNumMathWarps * MMA_M, "BLOCK_KV = warps × MMA_M");
     DG_STATIC_ASSERT(kHeadDim % MMA_K == 0 and kNumHeads % MMA_N == 0, "Alignment");
 
-    static constexpr uint32_t kSwizzleMode = 128;
+    // Must match the TMA descriptor's swizzle (host passes swizzle_mode = head_dim)
+    static constexpr uint32_t kSwizzleMode = kHeadDim;
     static constexpr uint32_t kSwizzleAlignment = kHeadDim * 8;
     static constexpr uint32_t kSMEMKBytes = kHeadDim;
 
diff --git a/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh b/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh
index f17c29ac7c..d35e90f1e4 100644
--- a/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh
+++ b/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh
@@ -53,7 +53,8 @@ void sm120_fp8_paged_mqa_logits(const uint32_t batch_size,
     static constexpr uint32_t kNumMathWarps = kNumMathThreads / 32;
     static constexpr uint32_t kWarpsPerGroup = BLOCK_KV / MMA_M;
     static constexpr uint32_t kNumGroups = kNumMathWarps / kWarpsPerGroup;
-    static constexpr uint32_t kSwizzleMode = 128;
+    // Must match the TMA descriptor's swizzle (host passes swizzle_mode = head_dim)
+    static constexpr uint32_t kSwizzleMode = kHeadDim;
     static constexpr uint32_t kSMEMKBytes = kHeadDim;
 
     DG_STATIC_ASSERT(kNumTMAThreads == 128, "Expected 128 TMA threads");
diff --git a/tests/test_attention.py b/tests/test_attention.py
index 1eb62e8b15..f5a8bd1f37 100644
--- a/tests/test_attention.py
+++ b/tests/test_attention.py
@@ -116,8 +116,12 @@ def enumerate_mqa_logits():
                     for seq_len, seq_len_kv, head_cfgs in (
                         *[(s, k, [(64, 128)]) for s in (128, 512, 2048, 4096) for k in (4096, 8192)],
                         *[(s, k, [(64, 128), (32, 128)]) for s in (510, 512) for k in (130560,)],
+                        # head_dim < 128 coverage (FP8-only): kernel swizzle must follow head_dim
+                        (512, 8192, [(64, 64), (64, 32)]),
                     ):
                         for num_heads, head_dim in head_cfgs:
+                            if is_fp4 and head_dim != 128:
+                                continue
                             for disable_cp in (False, True):
                                 yield is_fp4, logits_dtype, compressed_logits, clean_logits, seq_len, seq_len_kv, num_heads, head_dim, disable_cp