Skip to content

frost(sdpa): FP8/MXFP8 THD declared-stride support #537

Description

@vedaanta

Follow-up to #526 (native THD stride support in the f16 fwd kernels).

  • The SM120 per-tensor FP8 THD path (Add SM120 per-tensor FP8 (e4m3) SDPA-forward engine #509) keeps the packed contract: _thd_check_strides_packed declines non-packed declarations (AGENTS Hard Rule 2 — decline, never adapt). Extending native declared-stride support there needs the FP8 kernel's compile() to take per-tensor strides like the f16 kernels, with the 16-byte TMA global-stride rule evaluated at 1 byte/elem (strides x16 elements).
  • SM100 FP8/MXFP8 have no THD path wired at all yet; when added, they must adopt the declared-stride serve-or-decline contract from day one.
  • The fp8/mxfp8 test harness (test/python/sdpa/mxfp8.py etc.) allocates assuming packed strides, which is why test(sdpa): fuzz per-tensor ragged token-stride gaps in test_mhas_v2 #516's with_ragged_token_gap fuzz knob is enabled only in the fp16 ragged sweeps — the harness needs the same configured-stride allocation + offset generalization fp16.py got.

🤖 Generated with Claude Code

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions