Skip to content

SDPA DSL fp8/mxfp8 engines drop descale_s and scale_s, so non-unit values are silently wrong #520

Description

@YangXu1990uiuc

graph.sdpa_fp8 declares six scale operands. The DSL SDPA-forward engines
forward four of them and drop descale_s and scale_s before the adapter is
reached, so a graph supplying non-unit values builds, runs, and returns a
wrong result with no error.

Where

python/cudnn/sdpa/fwd/engines.py, the if facts.is_mxfp8 or facts.is_fp8:
block that assembles execute_kwargs, passes descale_q, descale_k,
descale_v, scale_o, the SF tensors and the amax buffers. descale_s and
scale_s are not there, and SdpaGraphFacts does not carry them either. Both
SdpaFwdDslSm100._execute_fp8 and SdpaFwdDslSm120._execute_fp8 therefore
compute:

scale_softmax_log2 = scale_val * descale_q * descale_k * log2(e)
o_scale_fused      = descale_v * scale_o

which is the correct math only when descale_s == scale_s == 1. Nothing
checks that. Affects the SM100 fp8 cell, the SM100 MXFP8 cell, and the SM120
fp8 cell.

Why it is silent

The kernels keep S in fp32 registers between QK and PV and quantize P with a
plain cvt.rn.satfinite.e4m3x2, i.e. an implicit scale_s of 1. The node
convention is that scale_s maps the post-softmax P into the e4m3 range and
descale_s undoes it on the PV output — precisely so P does not waste the
format's range. A caller following that convention gets O scaled by
1/scale_s and no diagnostic.

Every in-tree test passes 1.0 for both, which is why this has not shown up.

Two ways to close it

  • Fold. scale_s multiplies P, and P is exp2(...), so it can ride as
    +log2(scale_s) in the exponent addend for free; descale_s then folds into
    o_scale_fused. Care is needed because row_sum accumulates the unscaled
    P, so the two must not both pick up the factor. This needs a reference that
    exercises non-unit values before it can be trusted.
  • Reject. Forward both and raise at execute when either differs from 1.
    The fp8 paths already read four scalars with .item(), so two more cost
    nothing. This converts a wrong answer into a clear error, which is the actual
    defect.

Either one touches the execute signature shared by the SM100 and SM120 cells,
so it wants a change that can be exercised on both parts.

Found while addressing review on #509; deliberately not fixed there, since that
PR is an SM120 kernel change and this is neither SM120-specific nor kernel-level.

note to self: claude::304e9e55-1db7-4285-967f-001cb21032f3 — "审计 Frost Python DSL 引擎调用流程"
cwd /home/scratch.yanxu_libs/cudnn_frontend · workspace /home/scratch.yanxu_libs/fe_sm120

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions