[Perf][Kernel] Vectorized flat abs-max for dynamic per-tensor FP8 quantization (#58194)

Signed-off-by: Monishver Chandrasekaran <monishverchandrasekaran@gmail.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Co-authored-by: Misha Goin <mgoin64@gmail.com>

tests/kernels/quantization/test_fp8_quant.py

@@ -108,6 +108,36 @@ def test_dynamic_per_tensor_fp8_quant(
     opcheck_fp8_quant(ops_out, x)
 
 
+@pytest.mark.parametrize("pad", [8, 16])
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("seed", SEEDS)
+@torch.inference_mode()
+def test_dynamic_per_tensor_fp8_quant_strided_input(
+    pad: int, dtype: torch.dtype, seed: int
+) -> None:
+    """Non-contiguous rows take the strided abs-max kernel; contiguous rows take
+    the flat one. Both must produce exactly the same scale and bytes.
+
+    pad=8 (16 B for bf16) leaves rows misaligned for 32 B vector loads, pad=16
+    keeps them aligned but non-contiguous.
+    """
+    set_random_seed(seed)
+    num_tokens, hidden_size = 64, 4096
+
+    full = torch.rand(num_tokens, hidden_size + pad, dtype=dtype, device="cuda")
+    x = full[:, :hidden_size]
+    assert not x.is_contiguous()
+
+    ops_out, ops_scale = ops.scaled_fp8_quant(x)
+    flat_out, flat_scale = ops.scaled_fp8_quant(x.contiguous())
+    assert torch.equal(ops_scale, flat_scale)
+    assert torch.equal(ops_out.view(torch.uint8), flat_out.view(torch.uint8))
+
+    ref_out, ref_scale = ref_dynamic_per_tensor_fp8_quant(x)
+    assert torch.equal(ref_scale, ops_scale)
+    assert torch.equal(ref_out.view(torch.uint8), ops_out.view(torch.uint8))
+
+
 # Regression test for a case with large activations where an int32 index cannot
 # represent the number of elements.
 @torch.inference_mode()