[None][feat] skip-softmax on SM120: TMA-load + sync-MMA warp-specialized context FMHA for sm_120/sm_121 - #15163
Merged
Merged
Conversation
…el for sm_120/sm_121 Add the halfspec context FMHA kernel for the sm_120 family (sm_120/sm_121): a dedicated producer warp drives Q/K/V loads with TMA (cp.async.bulk.tensor + cuTensorMapEncodeTiled descriptors) into granular smem buffers, while the consumer warps run a BMM1 + softmax + per-warp skip-softmax + BMM2 body on mma.sync. sm_120/sm_121 have no wgmma.async, so only the load side of the Hopper warp-specialization recipe is ported -- hence 'halfspec'. The kernel reuses the existing LDGSTS Smem_tile_* types (their XOR swizzle is byte-identical to the TMA 128B hardware swizzle) and re-tiles V into 64-wide DV chunks so V smem rows stay 128 bytes (the only layout a TMA swizzle mode can fill). This adds an ENABLE_SKIP_SOFTMAX template parameter to the shared Kernel_traits_ so the halfspec traits can thread the skip-softmax knob through. See cpp/kernels/fmha_v2/src/fmha/warpspec_sm120/README.md for the full design rationale. Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
…hind use_halfspec_fmha
Compile the halfspec translation unit into the _context_attention_kernels_120
CMake target only (it uses sm_120-only TMA + sync-MMA), and add the
TLLM_ENABLE_HALFSPEC_SM120 guard so the all-architecture dispatch TU references
the run_halfspec_* bridge symbols only on builds that include sm_120.
Plumb a use_halfspec_fmha opt-in flag end to end: the PyTorch attention op
(nanobind) -> AttentionOp -> MHARunnerParams -> Launch_params. When set and the
config matches (sm_120/sm_121, BF16 in/out, causal, head_dim == head_dim_v in
{128, 256}, PACKED_QKV), FusedMultiHeadAttentionXMMAKernelV2::run dispatches to
the halfspec kernel; the flag is a no-op everywhere else. The TrtllmAttention
backend opts in whenever skip-softmax is configured, since halfspec is the
sm_120/sm_121 kernel that implements skip-softmax.
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
…d prefill The halfspec (TMA-load + sync-MMA warp-specialized) sm_120/sm_121 context FMHA issued its Q/K/V TMA loads with request-local sequence coordinates and no per-request token offset. Because the Q/K/V descriptors span the whole packed [total_tokens, H, D] buffer with a single base pointer, every batch element re-read request 0's tokens: single-request prefill was correct, but for batch size B > 1 all B requests returned request 0's result. Add the cumulative per-request token offset (binfo.sum_s == cu_q_seqlens[bidb]) to the Q/K/V seq coordinates in DMA::run. The KV offset reuses sum_s rather than sum_s_kv: halfspec is PACKED_QKV-only, so K/V share each request's token range with Q, and cu_kv_seqlens is null on the self-attention path (sum_s_kv would dereference it unconditionally in the Single_cta ctor, faulting). Add tests/unittest/_torch/attention/test_halfspec_sm120.py, which forces the halfspec kernel via a skip-softmax config (tiny threshold => no actual skipping, full softmax) and checks single- and multi-request batches against both an fp32 causal reference and the default TRTLLM context kernel. The multi-request cases are the regression guard for this fix. Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
Trim the verbose comment block on the per-request TMA offset and the test's module docstring / inline comments. No functional change. Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
Contributor
📝 WalkthroughWalkthroughThis PR introduces a warp-specialized "halfspec" fused multi-head flash attention kernel optimized for NVIDIA SM_120/SM_121. The implementation adds skip-softmax parameter support to kernel traits, defines SM120-specific producer-DMA and consumer-compute stages synchronized via mbarriers, integrates runtime dispatch conditioned on SM version and data format, and exposes the feature through Python binding and config flags. ChangesSM120 Halfspec Flash Attention Implementation
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (4)
cpp/kernels/fmha_v2/src/fmha/warpspec_sm120/dma_sync_mma.h (2)
315-316: 💤 Low valueFP16 support may need format distinction.
When
ELEMENT_BYTES == 2, this unconditionally usesCU_TENSOR_MAP_DATA_TYPE_BFLOAT16. For FP16 inputs, this should technically beCU_TENSOR_MAP_DATA_TYPE_FLOAT16. However, sinceCU_TENSOR_MAP_FLOAT_OOB_FILL_NONEis used (line 333), the format field only affects out-of-bounds fill behavior which is disabled here. The TMA moves raw bytes regardless of the format, so this works correctly for both BF16 and FP16 data in practice.Consider adding a comment clarifying this or accepting an explicit data type template parameter if FP16 support becomes a priority.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/kernels/fmha_v2/src/fmha/warpspec_sm120/dma_sync_mma.h` around lines 315 - 316, The conditional that sets CUtensorMapDataType based solely on Kernel_traits::ELEMENT_BYTES currently maps 2-byte elements to CU_TENSOR_MAP_DATA_TYPE_BFLOAT16 which is incorrect for FP16; update the code in dma_sync_mma.h (around the CUtensorMapDataType fmt assignment) to either accept an explicit data-type template parameter for true FP16 vs BF16 selection or, at minimum, add a clear comment explaining that because CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE is used the format only affects out-of-bounds fill and TMA moves raw bytes so the current mapping is intentional and works for both BF16 and FP16; reference Kernel_traits::ELEMENT_BYTES, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, CU_TENSOR_MAP_DATA_TYPE_FLOAT32 and CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE in the comment or when introducing the explicit type parameter.
334-343: 💤 Low valueError reporting doesn't prevent undefined behavior.
When
cuTensorMapEncodeTiledfails, this prints an error but continues execution. The kernel will likely crash or produce garbage with an uninitialized descriptor. For debug builds this is acceptable, but consider adding an assertion or a way to signal failure to the caller.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/kernels/fmha_v2/src/fmha/warpspec_sm120/dma_sync_mma.h` around lines 334 - 343, The current error handling after cuTensorMapEncodeTiled (checking res != CUDA_SUCCESS and only printing) can leave the tensor descriptor uninitialized and lead to UB; update the handling by either asserting/aborting on failure or propagating the error to the caller: inside the res != CUDA_SUCCESS branch (where err is retrieved), call assert(false) or abort() for debug builds, and for release builds set/return a failure code (e.g., return res or a bool/error enum) so the caller can bail out; ensure the chosen approach consistently prevents further use of the (uninitialized) descriptor created around the cuTensorMapEncodeTiled call.cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp (1)
509-521: ⚡ Quick winReuse the existing
isSm120fvariable instead of redeclaring it.Line 509 redeclares
isSm120f, which shadows the identical variable declared at Line 362 in the same function. While legal, this makes the code harder to maintain. Simply reuse the outerisSm120fand remove the redundant declaration here.♻️ Suggested simplification
mLaunchParams.enableSkipSoftmax = false; if (runnerParams.skipSoftmaxThresholdScaleFactor > 0) { - bool const isSm120f = (mSM == kSM_120 || mSM == kSM_121); bool const hopperWarpspec = isSm90 && mLaunchParams.warp_specialization; // The halfspec kernel is the only sm_120 / sm_121 FMHA that implements // skip-softmax, so skip-softmax there is only permitted with halfspec. bool const sm120Halfspec = isSm120f && mLaunchParams.useHalfspecFmha;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp` around lines 509 - 521, The local variable isSm120f is being redeclared in the block starting at the shown snippet; remove the redundant declaration and reuse the previously declared isSm120f (the one referencing mSM and kSM_120 / kSM_121) so the condition uses the outer isSm120f rather than shadowing it; update the condition that computes sm120Halfspec (and any subsequent logic that references isSm120f) to reference the existing isSm120f variable in fmhaRunner.cpp.tests/unittest/_torch/attention/test_halfspec_sm120.py (1)
35-35: ⚡ Quick winPrefer built-in
listovertyping.Listfor Python 3.10+ compatibility.Per coding guidelines, use built-in types (
list,dict,tuple) instead of legacytyping.List,typing.Dict,typing.Tuplein annotations. Since this codebase targets Python 3.10+, you can removefrom typing import Listand uselist[int]directly in the function signature at line 58.♻️ Suggested change
-from typing import ListThen update line 58:
- seq_lens: List[int], + seq_lens: list[int],🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/attention/test_halfspec_sm120.py` at line 35, Remove the legacy typing import and replace usages with built-in generics: delete the line "from typing import List" and update any function/type annotations that use "List" (notably the test function signature at line 58) to use the built-in form "list[int]" (or the appropriate element type), e.g., change "List[int]" to "list[int]" to conform to Python 3.10+ typing conventions.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/halfspec_sm120/fused_multihead_flash_attention_ws_sm120.cu`:
- Around line 163-165: The calls to ::launch_halfspec (e.g.,
::launch_halfspec<::fmha_halfspec::Halfspec_ktraits<256>>) currently ignore the
returned cuda error code; capture the return value (cudaError_t) from each
launch and propagate it instead of discarding it so that errors aren't masked by
the later cudaGetLastError() return in this function; update the function to
return the launch result (or if non-success, return that error immediately) for
both occurrences around the calls to ::launch_halfspec.
In `@tests/unittest/_torch/attention/test_halfspec_sm120.py`:
- Around line 129-132: Add a return type annotation to the test function
declaration for test_halfspec_context_matches_reference: update its signature to
include "-> None" so the function is explicitly typed as returning nothing,
ensuring compliance with the test suite coding guidelines and type checkers.
- Line 129: Update the pytest parameterization for the head_dim parameter so the
test covers both supported sizes: change the decorator argument from "head_dim",
[128] to "head_dim", [128, 256] (keeping the existing ids=lambda d:
f"head_dim_{d}") so the halfspec kernel is exercised for head_dim ∈ {128, 256}
in the test that uses the head_dim fixture/parameter.
- Around line 81-111: Add a return type annotation to _run_context indicating it
returns a tuple of (the attention layer instance, the Metadata object, and the
output tensor); e.g. annotate the function as returning Tuple[..., torch.Tensor]
and import/use typing.Tuple and typing.Any or a forward-reference to
AttentionCls.Metadata so the signature reflects: (layer instance from
AttentionCls, AttentionCls.Metadata, torch.Tensor) to make the return type
explicit for callers and linters.
---
Nitpick comments:
In `@cpp/kernels/fmha_v2/src/fmha/warpspec_sm120/dma_sync_mma.h`:
- Around line 315-316: The conditional that sets CUtensorMapDataType based
solely on Kernel_traits::ELEMENT_BYTES currently maps 2-byte elements to
CU_TENSOR_MAP_DATA_TYPE_BFLOAT16 which is incorrect for FP16; update the code in
dma_sync_mma.h (around the CUtensorMapDataType fmt assignment) to either accept
an explicit data-type template parameter for true FP16 vs BF16 selection or, at
minimum, add a clear comment explaining that because
CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE is used the format only affects out-of-bounds
fill and TMA moves raw bytes so the current mapping is intentional and works for
both BF16 and FP16; reference Kernel_traits::ELEMENT_BYTES,
CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, CU_TENSOR_MAP_DATA_TYPE_FLOAT32 and
CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE in the comment or when introducing the
explicit type parameter.
- Around line 334-343: The current error handling after cuTensorMapEncodeTiled
(checking res != CUDA_SUCCESS and only printing) can leave the tensor descriptor
uninitialized and lead to UB; update the handling by either asserting/aborting
on failure or propagating the error to the caller: inside the res !=
CUDA_SUCCESS branch (where err is retrieved), call assert(false) or abort() for
debug builds, and for release builds set/return a failure code (e.g., return res
or a bool/error enum) so the caller can bail out; ensure the chosen approach
consistently prevents further use of the (uninitialized) descriptor created
around the cuTensorMapEncodeTiled call.
In `@cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp`:
- Around line 509-521: The local variable isSm120f is being redeclared in the
block starting at the shown snippet; remove the redundant declaration and reuse
the previously declared isSm120f (the one referencing mSM and kSM_120 / kSM_121)
so the condition uses the outer isSm120f rather than shadowing it; update the
condition that computes sm120Halfspec (and any subsequent logic that references
isSm120f) to reference the existing isSm120f variable in fmhaRunner.cpp.
In `@tests/unittest/_torch/attention/test_halfspec_sm120.py`:
- Line 35: Remove the legacy typing import and replace usages with built-in
generics: delete the line "from typing import List" and update any function/type
annotations that use "List" (notably the test function signature at line 58) to
use the built-in form "list[int]" (or the appropriate element type), e.g.,
change "List[int]" to "list[int]" to conform to Python 3.10+ typing conventions.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 472f46e7-92ed-42ad-98e0-ba0a5dda1475
📒 Files selected for processing (18)
cpp/kernels/fmha_v2/src/fmha/kernel_traits.hcpp/kernels/fmha_v2/src/fmha/warpspec_sm120/README.mdcpp/kernels/fmha_v2/src/fmha/warpspec_sm120/compute_sync_mma.hcpp/kernels/fmha_v2/src/fmha/warpspec_sm120/dma_sync_mma.hcpp/kernels/fmha_v2/src/fmha/warpspec_sm120/kernel_traits.hcpp/kernels/fmha_v2/src/fused_multihead_flash_attention_kernel_ws_sm120.hcpp/tensorrt_llm/common/attentionOp.cppcpp/tensorrt_llm/common/attentionOp.hcpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/CMakeLists.txtcpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cppcpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.hcpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_v2.cppcpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/halfspec_sm120/fused_multihead_flash_attention_ws_sm120.cucpp/tensorrt_llm/nanobind/thop/bindings.cppcpp/tensorrt_llm/thop/attentionOp.cppcpp/tensorrt_llm/thop/attentionOp.htensorrt_llm/_torch/attention_backend/trtllm.pytests/unittest/_torch/attention/test_halfspec_sm120.py
The "halfspec" codename did not convey the kernel's purpose. Rename it to skip_softmax to match the existing skip-softmax feature naming (SkipSoftmaxAttentionConfig, skip_softmax_threshold_scale_factor, enableSkipSoftmax). Mechanical rename across identifiers, the TLLM_ENABLE_SKIP_SOFTMAX_SM120 build define, the use_skip_softmax_fmha dispatch flag, namespaces/traits, the directory (skip_softmax_sm120/) and the test file (test_skip_softmax_sm120.py). README/comment prose updated where the old "half of the warp-spec recipe" etymology no longer applied. No behavior change. Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
dcampora
changed the title
yuxianq
reviewed
Jun 11, 2026
yuxianq
reviewed
Jun 11, 2026
yuxianq
reviewed
Jun 11, 2026
yuxianq
reviewed
Jun 11, 2026
dcampora
force-pushed
the
dcampora/halfspec-sm120
branch
from
June 12, 2026 07:18
f5bdeb5 to
417dc5f
Compare
…fault context kernel
The TMA-load + sync-MMA warp-specialized context FMHA is now the default
context attention path for sm_120 / sm_121 rather than an opt-in.
Skip-softmax is driven by the prefill threshold alone -- there is no separate
enable flag. The sm_120 / sm_121 bridges read
params.skip_softmax_threshold_scale_factor directly to select the
ENABLE_SKIP_SOFTMAX kernel variant: the false variant is a plain full-softmax
prefill (no skip-check overhead) and runs by default; the true variant runs when
a threshold is set. If a threshold is set but no skip-capable kernel matches the
config, skipping is simply not enabled and the request runs full softmax (no
error).
- Drop the use_skip_softmax_fmha / useSkipSoftmaxFmha / mUseSkipSoftmaxFmha opt-in
end-to-end: the attention op + nanobind binding, MHARunnerParams / Launch_params,
AttentionOp (incl. its data() identity tuple), fmhaRunner, and the Python
TrtllmAttention property + call site.
- Launch_params::enableSkipSoftmax is now only the Hopper cubin-selection bit
(threshold > 0 && warp_specialization && flash_attention); on sm_120 / sm_121 it
stays false and the threshold drives the bridge directly.
- Dispatch every matching sm_120 / sm_121 prefill (BF16 in/out, causal, packed QKV,
head_dim == head_dim_v in {128, 256}) to the kernel, with guards that fall back to
the cubin path for features it does not implement: alibi (params.has_alibi),
logit softcapping, sage attention, returning softmax stats
(params.softmax_stats_ptr), and the interleaved layout.
- Propagate the launch_skip_softmax CUDA error from the bridges (it was dropped,
which could clear/hide launch failures).
- Abort on a cuTensorMapEncodeTiled failure instead of continuing with an
unencoded TMA descriptor (the prior print-and-continue left undefined behavior),
and document the BF16-only descriptor data-type mapping.
- Remove the unused RING_DEPTH kernel-traits knob.
- Test: exercise both kernel variants (no-skip default + 1e-30 skip) against the
fp32 reference across head_dim in {128, 256}; modernize annotations (return
types and built-in list over typing.List).
TLLM_ENABLE_SKIP_SOFTMAX_SM120 is retained as a link-availability guard: it lets the
all-architecture dispatch TU reference the sm_120-only bridge symbols without an
undefined-symbol error on builds that exclude sm_120. It is not a feature gate.
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
yuxianq
approved these changes
Jun 12, 2026
dcampora
force-pushed
the
dcampora/halfspec-sm120
branch
from
June 12, 2026 12:28
417dc5f to
fffc18c
Compare
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #53887 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #53887 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #53951 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #53951 [ run ] completed with state |
1 task
xinhe-nv
pushed a commit
to tensorrt-cicd/TensorRT-LLM
that referenced
this pull request
Jun 23, 2026
…zed context FMHA for sm_120/sm_121 (NVIDIA#15163) Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> Signed-off-by: GitLab CI Bot <gitlab-ci@nvidia.com>
xinhe-nv
pushed a commit
to tensorrt-cicd/TensorRT-LLM
that referenced
this pull request
Jun 24, 2026
…zed context FMHA for sm_120/sm_121 (NVIDIA#15163) Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> Signed-off-by: GitLab CI Bot <gitlab-ci@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Adds skip_softmax, a warp-specialized context (prefill) FMHA for the
sm_120 / sm_121 (consumer Blackwell) family. It is the sm_120 / sm_121 attention
path that implements the per-warp skip-softmax optimization (hence the name).
Only half of the Hopper warp-specialization recipe ports to consumer Blackwell:
the TMA-driven async loads survive, but async MMA does not (no
wgmma.async),so a single producer warp issues
cp.async.bulk.tensorloads for Q/K/V intogranular shared-memory buffers via driver-API
CUtensorMapdescriptors, whilethe consumer warps run BMM1 + softmax (+ per-warp skip-softmax) + BMM2 on
mma.sync, synchronized by an mbarrier producer/consumer handshake.It is opt-in and activates when a
SkipSoftmaxAttentionConfigis attached tothe attention layer (
TrtllmAttention.use_skip_softmax_fmha). The C++ runnerrestricts dispatch to the supported shapes — BF16 in/out,
head_dim == head_dim_vin {128, 256}, causal mask, packed QKV — and ignoresthe flag everywhere else, so it is a no-op on other hardware and configs.
What's included:
cpp/kernels/fmha_v2/src/fmha/warpspec_sm120/):kernel_traits.h,dma_sync_mma.h(producer / TMA),compute_sync_mma.h(consumer), plus theentry header
fused_multihead_flash_attention_kernel_ws_sm120.hand a designdoc
README.md. Compiled sm_120-only via the_context_attention_kernels_120CMake target, guarded by
TLLM_ENABLE_SKIP_SOFTMAX_SM120.use_skip_softmax_fmhais plumbed from thethop.attentionop →AttentionOp→MHARunnerParams→Launch_params→FusedMultiHeadAttentionXMMAKernelV2::run, which routes the matching config tothe
run_skip_softmax_*bridges. Python property intensorrt_llm/_torch/attention_backend/trtllm.py.token offset (
cu_q_seqlens[bidb]) to the Q/K/V TMA sequence coordinates;without it, multi-request prefill batches all returned request 0's result.
Verified on an RTX PRO 6000 Blackwell (sm_120).
Test Coverage
tests/unittest/_torch/attention/test_skip_softmax_sm120.py(new): forces theskip_softmax kernel via a
SkipSoftmaxAttentionConfigwith a tinythreshold_scale_factor(no tiles skipped ⇒ full softmax) and checks single-and multi-request causal prefill against both an fp32 reference and the default
(non-skip_softmax) TRTLLM context kernel. Skips on non-sm_120/121 GPUs; the
multi-request cases are the regression guard for the per-request TMA offset.
tests/unittest/_torch/attention_backend/test_attention_op_sync.py(existing):guards the
use_skip_softmax_fmhakwarg plumbing across thethop.attentionboundary.
PR Checklist
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.