Unverified Commit 941364df authored by Charlene Yang's avatar Charlene Yang Committed by GitHub
Browse files

Fix docstring related to `t` in `thd` (#1111)



fix typos regarding t in thd
Signed-off-by: default avatarCharlene Yang <8636796+cyanguwa@users.noreply.github.com>
parent a326e351
......@@ -22,7 +22,7 @@ extern "C" {
/*! \enum NVTE_QKV_Layout
* \brief Memory layouts of QKV tensors.
* `S`, `B`, `H`, `D`, and `T` stand for sequence length, batch size, number of heads,
* head size, and the total number of sequences in a batch, i.e. `t = sum(s_i) for i = 0...b-1`.
* head size, and the total number of tokens in a batch, i.e. `t = sum(s_i) for i = 0...b-1`.
* `SBHD` and `BSHD`-based layouts are used when sequences in a batch are of equal length
* or padded to the same length, and `THD`-based layouts are used when sequences have
* different lengths in a batch.
......
......@@ -3122,7 +3122,7 @@ def get_qkv_layout(
qkv_format: str, default = `sbhd`
Dimension format for `q`, `k` and `v`, {`sbhd`, `bshd`, `thd`}. `s` stands for
the sequence length dimension, `b` batch size, `h` the number of attention heads,
`d` head size, and `t` the total number of sequences in a batch, i.e.
`d` head size, and `t` the total number of tokens in a batch, i.e.
`t = sum(s_i) for i = 0...b-1`.
Returns
......@@ -5232,7 +5232,7 @@ class DotProductAttention(TransformerEngineBaseModule):
qkv_format: str, default = `sbhd`
dimension format for `query_layer`, `key_layer` and `value_layer`,
{`sbhd`, `bshd`, `thd`}. `s` stands for the sequence length, `b` batch size,
`h` the number of heads, `d` head size, and `t` the total number of sequences
`h` the number of heads, `d` head size, and `t` the total number of tokens
in a batch, with `t = sum(s_i), for i = 0...b-1`. `sbhd` and `bshd` formats
are used for when sequences in a batch are of equal length or padded to
equal length, and the `thd` format is used for when sequences in a batch
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment