You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
The THD lowering reconstructs each sequence's base row as prefix(lens) x token_stride and never reads the ragged-offset tensors the caller bound. fwd/engines.py states this outright:
THD lowerings assume FULLY-PACKED storage: the packed addressing is re-derived as prefix(lens) x token stride, and the graph's bound ragged-offset values are never read. TE-style padded THD (offsets from cu_seqlens_padded != cu_seqlens, gaps between sequences) is NOT served — and being runtime data, cannot be declined at plan time.
That last clause is the problem. A padded-THD caller is not rejected; the graph is claimed and silently mis-addressed.
Concrete failure
TE's DotProductAttention documents this layout directly. A batch of 3 sequences with identifier tokens inserted between them:
[a a a 1 b b 2 2 c c c c 3]
cu_seqlens = [0, 3, 5, 9]
cu_seqlens_padded = [0, 4, 8, 13]
Sequence 1 starts at row 4 (b). We would derive prefix(lens)[1] = 3 and start it at row 3 — the identifier token 1. Every sequence after the first is shifted, output is wrong, and nothing reports an error.
Memory safety is not at risk: the per-sequence offset enters as a TMA coordinate (tma_k(0, head, kv_row_base + ... + kv_seq_off, tma_batch)), which the hardware bounds-checks against GLOBAL_DIM, and #706 bounds capacity by the view's element span. So this is silent wrong results, not a fault.
Proposal: consume the bound offsets on device
The offsets are already device tensors in the variant pack, and the setup kernel already runs on device and writes the [seq_kv | cu_q | cu_k] metadata. It could read the caller's ragged-offset tensors instead of reconstructing offsets from lengths:
packed callers unaffected — bound offsets equal the derived ones by definition;
padded callers become correctly served rather than silently mis-served;
This converts "cannot be declined at plan time" into "does not need to be declined" — strictly better than any validation we could add, since validating a device-resident prefix sum would require exactly the D2H read #552 removed.
Element vs byte convention for the bound offsets, per tensor.
Per-tensor offsets: Q/K/V/O each carry their own ragged-offset tensor, so Q/O padding can differ from K/V padding. The current single cu_q / cu_k metadata pair cannot express that — the metadata layout would need to widen.
Lengths are still required for mask bounds and the dead-row path, so this adds device reads rather than replacing them.
Capacity: _thd_capacity currently reasons about a packed span; under padding it must cover the padded span.
Stats/LSE carry their own ragged offsets and would need the same treatment.
Padded THD makes #624 substantially worse. In packed THD an interior tile tail steps into the next sequence's real data (finite, harmless), so only the final sequence's tail is exposed — one tile per launch. Under padding every sequence's tail tile lands in its own pad slot, so the exposure becomes B tails. TE's identifier-token case keeps real embeddings there, but the context-parallel case pads to uniform 2*cp_size chunks with no initialization guarantee. Whatever fix #624 lands should be chosen with padded THD in mind.
Why it matters
TransformerEngine and Megatron-Core are the primary padded-THD producers (PackedSeqParams carries cu_seqlens_q_padded / cu_seqlens_kv_padded straight through), so this sits on the critical path for TE adopting the cuDNN Python API — and today the failure mode there would be silently wrong numbers rather than a clean decline.
Summary
The THD lowering reconstructs each sequence's base row as
prefix(lens) x token_strideand never reads the ragged-offset tensors the caller bound.fwd/engines.pystates this outright:That last clause is the problem. A padded-THD caller is not rejected; the graph is claimed and silently mis-addressed.
Concrete failure
TE's
DotProductAttentiondocuments this layout directly. A batch of 3 sequences with identifier tokens inserted between them:Sequence 1 starts at row 4 (
b). We would deriveprefix(lens)[1] = 3and start it at row 3 — the identifier token1. Every sequence after the first is shifted, output is wrong, and nothing reports an error.Memory safety is not at risk: the per-sequence offset enters as a TMA coordinate (
tma_k(0, head, kv_row_base + ... + kv_seq_off, tma_batch)), which the hardware bounds-checks againstGLOBAL_DIM, and #706 bounds capacity by the view's element span. So this is silent wrong results, not a fault.Proposal: consume the bound offsets on device
The offsets are already device tensors in the variant pack, and the setup kernel already runs on device and writes the
[seq_kv | cu_q | cu_k]metadata. It could read the caller's ragged-offset tensors instead of reconstructing offsets from lengths:This converts "cannot be declined at plan time" into "does not need to be declined" — strictly better than any validation we could add, since validating a device-resident prefix sum would require exactly the D2H read #552 removed.
What needs checking before calling it simple
CUDNN_ATTR_TENSOR_RAGGED_OFFSET_MULTIPLIER, 9.24+): bound values may be scaled, and the multiplier has to be applied on the device side.cu_q/cu_kmetadata pair cannot express that — the metadata layout would need to widen._thd_capacitycurrently reasons about a packed span; under padding it must cover the padded span.cu_k[B](see frost(sdpa): THD zero-host-read execute loads uninitialized KV capacity rows — NaN poisons P@V on the f16/fp8 SM100/SM120 rows #624); under padding that becomescu_k_padded[B].Interaction with #624
Padded THD makes #624 substantially worse. In packed THD an interior tile tail steps into the next sequence's real data (finite, harmless), so only the final sequence's tail is exposed — one tile per launch. Under padding every sequence's tail tile lands in its own pad slot, so the exposure becomes B tails. TE's identifier-token case keeps real embeddings there, but the context-parallel case pads to uniform
2*cp_sizechunks with no initialization guarantee. Whatever fix #624 lands should be chosen with padded THD in mind.Why it matters
TransformerEngine and Megatron-Core are the primary padded-THD producers (
PackedSeqParamscarriescu_seqlens_q_padded/cu_seqlens_kv_paddedstraight through), so this sits on the critical path for TE adopting the cuDNN Python API — and today the failure mode there would be silently wrong numbers rather than a clean decline.