Fix TMA descriptor error on H100 PCIe: pad LSE buffer width to multiple of 16 - #13
Open
vivien-cheng wants to merge 1 commit into
Open
Fix TMA descriptor error on H100 PCIe: pad LSE buffer width to multiple of 16#13vivien-cheng wants to merge 1 commit into
vivien-cheng wants to merge 1 commit into
Conversation
The kernel rounds the attention-partials axis up to a multiple of 16 (ROUNDED_MAX_ATTN_PARTIALS in attention_reduction.cu, and the sv_fl width in attn_lse_intermediates_t in llama.cuh) and TMA-loads the buffer at that padded width. The Python allocation used the raw SM count, so the global tensor was narrower than the TMA box. On an SM count not divisible by 4 this also makes the innermost stride not a multiple of 16 bytes, which cuTensorMapEncodeTiled rejects outright. Apply the same round-up in Python. The padding columns are never indexed.
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
1 participant
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.
Ran the demo on H100 PCIe (114 SMs), got stderr floods with "Error in vector TMA descriptor creation: invalid argument" and the attention reduction reads garbage.
The kernel pads the attention-partials axis to a multiple of 16 (
ROUNDED_MAX_ATTN_PARTIALS), but the Python allocation uses the raw SM count, so the buffer is narrower than the TMA box and the stride isn't 16-byte aligned.Fix: same round-up in Python. Padding columns are never indexed.
Tested on H100 PCIe: clean build, no TMA errors (673 tok/s).