Raise RuntimeError on unsupported head_dim/dtype in SM100 dense FMHA dispatch - #194
Raise RuntimeError on unsupported head_dim/dtype in SM100 dense FMHA dispatch#194Functionhx wants to merge 1 commit into
Conversation
…l (fwd+bwd) When an unsupported (head_dim_qk, head_dim_vo) pair is passed to the SM100 dense prefill forward or backward kernel, the code previously fell through to a std::cout message and silently returned without launching a kernel. Since the output tensors were pre-allocated with torch::empty(), callers received uninitialized (garbage) data instead of a clear error. Root cause: The head_dim dispatch in fmha_cutlass_fwd_sm100.cu and fmha_cutlass_bwd_sm100.cu only handled (192,128) and (128,128). Any other combination hit the else branch which only printed to stdout. Fix: Replace the std::cout fallthrough with TORCH_CHECK(false, ...) that raises a RuntimeError including the unsupported values and the list of supported combinations. Test Plan: Ran: pip install -e . -v --no-build-isolation (build succeeded) Fixes: deepseek-ai#185 Signed-off-by: Yuchen Fan <functionhx@gmail.com>
ff63bba to
11d11f1
Compare
|
For maintainers' reference: the head_dim half of this PR is the same change as #185 (open since May 30) — the TORCH_CHECK and its message text match #185's, with a suffix added. #185 also includes the regression test requested in its review (tests/test_fmha_sm100_dispatch.py); this PR contains no test, although the author's comment on #185 said it would include one. The genuinely new part here is the dtype branch (FLASH_MLA_ASSERT → TORCH_CHECK), which #185 intentionally left out of scope. Two process notes: "Fixes #185" references an open PR, not an issue. And this PR was opened at 18:07:26 UTC — 32 seconds after asking on #185 whether to open one (18:06:54 UTC), before anyone could reply. Happy to defer to the maintainers on how to combine the two. |
|
Closing — this duplicates #185 (open since May 30) which already has the same head_dim fix plus a regression test. The dtype change belongs as a separate, scoped PR. Apologies for the noise. |
2 participants
What Problem This Solves
Fixes #185: When unsupported
(head_dim_qk, head_dim_vo)values are passed to the SM100 dense prefill forward/backward kernels, the code prints a message to stdout and returns without launching a kernel. The output tensors are pre-allocated withtorch::empty(), so callers receive uninitialized/garbage data instead of an error.Fix
Replaced silent fallthroughs with
TORCH_CHECK(false, ...)in both forward and backward kernels:RuntimeErrorwith actual values and supported combinations(192, 128) and (128, 128)FLASH_MLA_ASSERT(false)→std::abort()(SIGABRT). Now raisesRuntimeErrorwith expected vs actual dtypestd::coutoutput removedEvidence
Two files changed:
fmha_cutlass_fwd_sm100.cuandfmha_cutlass_bwd_sm100.cu. Nostd::coutorFLASH_MLA_ASSERTremain.